diff --git a/stream.go b/stream.go index afe9199..b65f28a 100644 --- a/stream.go +++ b/stream.go @@ -67,14 +67,19 @@ func (s *Stream) Close() { close(s.ch) } -// OnClose 注册关闭回调;流已关闭时立即执行 +// OnClose 注册关闭回调;可在业务方法 return *Stream 之前同步调用。 +// 不得等待 Inject:Inject 只有业务方法返回后才发生,等待会形成循环依赖死锁。 +// 流已关闭时在锁外立即执行回调,避免回调重入 Stream 时自锁。 func (s *Stream) OnClose(cb func()) { - <-s.getReady() + if cb == nil { + return + } s.mu.Lock() - defer s.mu.Unlock() if s.closed { + s.mu.Unlock() cb() return } s.onClose = cb + s.mu.Unlock() } diff --git a/stream_onclose_test.go b/stream_onclose_test.go new file mode 100644 index 0000000..b5ad11a --- /dev/null +++ b/stream_onclose_test.go @@ -0,0 +1,46 @@ +package fun + +import ( + "sync/atomic" + "testing" + "time" +) + +// OnClose 必须允许业务方法在 return *Stream 之前同步注册;旧实现等待 Inject, +// 而 Inject 只有 invoke 在业务方法返回后才调用,形成精确死锁: +// +// business OnClose -> wait ready -> invoke wait business return -> never Inject. +func TestStreamOnCloseBeforeInjectDoesNotBlock(t *testing.T) { + st := &Stream{} + registered := make(chan struct{}) + var called atomic.Int32 + go func() { + st.OnClose(func() { called.Add(1) }) + close(registered) + }() + select { + case <-registered: + // expected: registration is independent from channel injection + case <-time.After(time.Second): + t.Fatal("OnClose blocked before Inject (business method would deadlock)") + } + + ch := make(chan any) + done := make(chan struct{}) + st.Inject(ch, done) + st.Close() + if got := called.Load(); got != 1 { + t.Fatalf("callback called %d times, want 1", got) + } +} + +func TestStreamOnCloseAfterAlreadyClosedRunsImmediately(t *testing.T) { + st := &Stream{} + st.Inject(make(chan any), make(chan struct{})) + st.Close() + var called atomic.Int32 + st.OnClose(func() { called.Add(1) }) + if got := called.Load(); got != 1 { + t.Fatalf("callback called %d times, want 1", got) + } +}