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

     1  package router
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/asynkron/protoactor-go/actor"
     7  )
     8  
     9  type RouterType int
    10  
    11  const (
    12  	GroupRouterType RouterType = iota
    13  	PoolRouterType
    14  )
    15  
    16  type RouterConfig interface {
    17  	RouterType() RouterType
    18  	OnStarted(context actor.Context, props *actor.Props, state State)
    19  	CreateRouterState() State
    20  }
    21  
    22  type GroupRouter struct {
    23  	Routees *actor.PIDSet
    24  }
    25  
    26  type PoolRouter struct {
    27  	PoolSize int
    28  }
    29  
    30  func (config *GroupRouter) OnStarted(context actor.Context, props *actor.Props, state State) {
    31  	config.Routees.ForEach(func(i int, pid *actor.PID) {
    32  		context.Watch(pid)
    33  	})
    34  	state.SetSender(context)
    35  	state.SetRoutees(config.Routees)
    36  }
    37  
    38  func (config *GroupRouter) RouterType() RouterType {
    39  	return GroupRouterType
    40  }
    41  
    42  func (config *PoolRouter) OnStarted(context actor.Context, props *actor.Props, state State) {
    43  	var routees actor.PIDSet
    44  	for i := 0; i < config.PoolSize; i++ {
    45  		routees.Add(context.Spawn(props))
    46  	}
    47  	state.SetSender(context)
    48  	state.SetRoutees(&routees)
    49  }
    50  
    51  func (config *PoolRouter) RouterType() RouterType {
    52  	return PoolRouterType
    53  }
    54  
    55  func spawner(config RouterConfig) actor.SpawnFunc {
    56  	return func(actorSystem *actor.ActorSystem, id string, props *actor.Props, parentContext actor.SpawnerContext) (*actor.PID, error) {
    57  		return spawn(actorSystem, id, config, props, parentContext)
    58  	}
    59  }
    60  
    61  func spawn(actorSystem *actor.ActorSystem, id string, config RouterConfig, props *actor.Props, parentContext actor.SpawnerContext) (*actor.PID, error) {
    62  	ref := &process{
    63  		actorSystem: actorSystem,
    64  	}
    65  	proxy, absent := actorSystem.ProcessRegistry.Add(ref, id)
    66  	if !absent {
    67  		return proxy, actor.ErrNameExists
    68  	}
    69  
    70  	pc := *props
    71  	pc.Configure(actor.WithSpawnFunc(nil))
    72  	ref.state = config.CreateRouterState()
    73  
    74  	if config.RouterType() == GroupRouterType {
    75  		wg := &sync.WaitGroup{}
    76  		wg.Add(1)
    77  		ref.router, _ = actor.DefaultSpawner(actorSystem, id+"/router", actor.PropsFromProducer(func() actor.Actor {
    78  			return &groupRouterActor{
    79  				props:  &pc,
    80  				config: config,
    81  				state:  ref.state,
    82  				wg:     wg,
    83  			}
    84  		}), parentContext)
    85  		wg.Wait() // wait for routerActor to start
    86  	} else {
    87  		wg := &sync.WaitGroup{}
    88  		wg.Add(1)
    89  		ref.router, _ = actor.DefaultSpawner(actorSystem, id+"/router", actor.PropsFromProducer(func() actor.Actor {
    90  			return &poolRouterActor{
    91  				props:  &pc,
    92  				config: config,
    93  				state:  ref.state,
    94  				wg:     wg,
    95  			}
    96  		}), parentContext)
    97  		wg.Wait() // wait for routerActor to start
    98  	}
    99  
   100  	ref.parent = parentContext.Self()
   101  	return proxy, nil
   102  }