Fix poison deadlocking on non-existing pids

actor/engine.go

@@ -225,24 +225,21 @@ func (e *Engine) sendPoisonPill(pid *PID, graceful bool, wg ...*sync.WaitGroup)
 	} else {
 		_wg = &sync.WaitGroup{}
 	}
-	_wg.Add(1)
-	proc := e.Registry.get(pid)
+	pill := poisonPill{
+		wg:       _wg,
+		graceful: graceful,
+	}
 	// deadletter - if we didn't find a process, we will broadcast a DeadletterEvent
-	if proc == nil {
+	if e.Registry.get(pid) == nil {
 		e.BroadcastEvent(DeadLetterEvent{
 			Target:  pid,
-			Message: poisonPill{_wg, graceful},
+			Message: pill,
 			Sender:  nil,
 		})
 		return _wg
 	}
-	pill := poisonPill{
-		wg:       _wg,
-		graceful: graceful,
-	}
-	if proc != nil {
-		e.SendLocal(pid, pill, nil)
-	}
+	_wg.Add(1)
+	e.SendLocal(pid, pill, nil)
 	return _wg
 }
 

actor/engine_test.go

@@ -284,7 +284,7 @@ func TestStop(t *testing.T) {
 
 func TestPoisonWaitGroup(t *testing.T) {
 	var (
-		wg = sync.WaitGroup{}
+		wg = &sync.WaitGroup{}
 		x  = int32(0)
 	)
 	e, err := NewEngine(NewEngineConfig())
@@ -303,6 +303,21 @@ func TestPoisonWaitGroup(t *testing.T) {
 
 	e.Poison(pid).Wait()
 	assert.Equal(t, int32(1), atomic.LoadInt32(&x))
+
+	// validate poisoning non exiting pid does not deadlock
+	wg = e.Poison(NewPID(LocalLookupAddr, "non-existing"))
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		wg.Wait()
+	}()
+	select {
+	case <-done:
+	case <-time.After(20 * time.Millisecond):
+		t.Error("poison waitGroup deadlocked")
+	}
+	// ... or panic
+	e.Poison(nil).Wait()
 }
 
 func TestPoison(t *testing.T) {
@@ -327,7 +342,6 @@ func TestPoison(t *testing.T) {
 		// When a process is poisoned it should be removed from the registry.
 		// Hence, we should get NIL when we try to get it.
 		assert.Nil(t, e.Registry.get(pid))
-
 	}
 }
 

actor/registry.go

@@ -40,6 +40,9 @@ func (r *Registry) Remove(pid *PID) {
 // If it doesn't exist, nil is returned so the caller must check for that
 // and direct the message to the deadletter processer instead.
 func (r *Registry) get(pid *PID) Processer {
+	if pid == nil {
+		return nil
+	}
 	r.mu.RLock()
 	defer r.mu.RUnlock()
 	if proc, ok := r.lookup[pid.ID]; ok {