diff --git a/pool.go b/pool.go index 0ec096f..2de26c3 100644 --- a/pool.go +++ b/pool.go @@ -368,11 +368,11 @@ func (p *Pool) Connect(id string, opts ...ConnectOption) error { InboxCounter: p.inboxCounter, Retire: func(err error) { p.mu.Lock() - defer p.mu.Unlock() if p.closed { return } delete(p.peers, id) + p.mu.Unlock() p.events <- PoolEvent{ID: id, Kind: EventRetired, Err: err, At: time.Now()} }, } diff --git a/worker_test.go b/worker_test.go index 856d9b7..8722259 100644 --- a/worker_test.go +++ b/worker_test.go @@ -189,7 +189,7 @@ func TestWorkerSession(t *testing.T) { } }, "expected worker to exit after Stop") - assert.False(t, retired.Load(), "expected Retire not called when context cancelled") + assert.False(t, retired.Load(), "expected Retire not called") }) t.Run("Retire not called on successful connection", func(t *testing.T) { @@ -200,7 +200,8 @@ func TestWorkerSession(t *testing.T) { w := makeWorker(t, mockSocket, nil, ctx, cancel) _, events, pool := makeWorkerContext(t) - pool.Retire = func(_ error) { panic("Retire must not be called on successful connection") } + retired := atomic.Bool{} + pool.Retire = func(_ error) { retired.Store(true) } var wg sync.WaitGroup wg.Go(func() { w.Start(pool) }) @@ -213,6 +214,8 @@ func TestWorkerSession(t *testing.T) { return false } }, "expected EventConnected without Retire being called") + + assert.False(t, retired.Load(), "expected Retire not called") }) t.Run("Stop before connection established - exits cleanly, no events", func(t *testing.T) {