github.com/asynkron/protoactor-go@v0.0.0-20240308120642-ef91a6abee75/router/routeractor_pool.go (about)

     1  package router
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/asynkron/protoactor-go/actor"
     8  )
     9  
    10  type poolRouterActor struct {
    11  	props  *actor.Props
    12  	config RouterConfig
    13  	state  State
    14  	wg     *sync.WaitGroup
    15  }
    16  
    17  func (a *poolRouterActor) Receive(context actor.Context) {
    18  	switch m := context.Message().(type) {
    19  	case *actor.Started:
    20  		a.config.OnStarted(context, a.props, a.state)
    21  		a.wg.Done()
    22  
    23  	case *AddRoutee:
    24  		r := a.state.GetRoutees()
    25  		if r.Contains(m.PID) {
    26  			return
    27  		}
    28  		context.Watch(m.PID)
    29  		r.Add(m.PID)
    30  		a.state.SetRoutees(r)
    31  
    32  	case *RemoveRoutee:
    33  		r := a.state.GetRoutees()
    34  		if !r.Contains(m.PID) {
    35  			return
    36  		}
    37  
    38  		context.Unwatch(m.PID)
    39  		r.Remove(m.PID)
    40  		a.state.SetRoutees(r)
    41  		// sleep for 1ms before sending the poison pill
    42  		// This is to give some time to the routee actor receive all
    43  		// the messages. Specially due to the synchronization conditions in
    44  		// consistent hash router, where a copy of hmc can be obtained before
    45  		// the update and cause messages routed to a dead routee if there is no
    46  		// delay. This is a best effort approach and 1ms seems to be acceptable
    47  		// in terms of both delay it cause to the router actor and the time it
    48  		// provides for the routee to receive messages before it dies.
    49  		time.Sleep(time.Millisecond * 1)
    50  		context.Send(m.PID, &actor.PoisonPill{})
    51  
    52  	case *BroadcastMessage:
    53  		msg := m.Message
    54  		sender := context.Sender()
    55  		a.state.GetRoutees().ForEach(func(i int, pid *actor.PID) {
    56  			context.RequestWithCustomSender(pid, msg, sender)
    57  		})
    58  
    59  	case *GetRoutees:
    60  		r := a.state.GetRoutees()
    61  		routees := make([]*actor.PID, r.Len())
    62  		r.ForEach(func(i int, pid *actor.PID) {
    63  			routees[i] = pid
    64  		})
    65  
    66  		context.Respond(&Routees{PIDs: routees})
    67  	case *actor.Terminated:
    68  		r := a.state.GetRoutees()
    69  		if r.Remove(m.Who) {
    70  			a.state.SetRoutees(r)
    71  		}
    72  	}
    73  }