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

     1  package router
     2  
     3  import (
     4  	"log"
     5  
     6  	"github.com/asynkron/protoactor-go/actor"
     7  	"github.com/serialx/hashring"
     8  )
     9  
    10  type Hasher interface {
    11  	Hash() string
    12  }
    13  
    14  type consistentHashGroupRouter struct {
    15  	GroupRouter
    16  }
    17  
    18  type consistentHashPoolRouter struct {
    19  	PoolRouter
    20  }
    21  
    22  type hashmapContainer struct {
    23  	hashring  *hashring.HashRing
    24  	routeeMap map[string]*actor.PID
    25  }
    26  type consistentHashRouterState struct {
    27  	hmc    *hashmapContainer
    28  	sender actor.SenderContext
    29  }
    30  
    31  func (state *consistentHashRouterState) SetSender(sender actor.SenderContext) {
    32  	state.sender = sender
    33  }
    34  
    35  func (state *consistentHashRouterState) SetRoutees(routees *actor.PIDSet) {
    36  	// lookup from node name to PID
    37  	hmc := hashmapContainer{}
    38  	hmc.routeeMap = make(map[string]*actor.PID)
    39  	nodes := make([]string, routees.Len())
    40  	routees.ForEach(func(i int, pid *actor.PID) {
    41  		nodeName := pid.Address + "@" + pid.Id
    42  		nodes[i] = nodeName
    43  		hmc.routeeMap[nodeName] = pid
    44  	})
    45  	// initialize hashring for mapping message keys to node names
    46  	hmc.hashring = hashring.New(nodes)
    47  	state.hmc = &hmc
    48  }
    49  
    50  func (state *consistentHashRouterState) GetRoutees() *actor.PIDSet {
    51  	var routees actor.PIDSet
    52  	hmc := state.hmc
    53  	for _, v := range hmc.routeeMap {
    54  		routees.Add(v)
    55  	}
    56  	return &routees
    57  }
    58  
    59  func (state *consistentHashRouterState) RouteMessage(message interface{}) {
    60  	_, uwpMsg, _ := actor.UnwrapEnvelope(message)
    61  	switch msg := uwpMsg.(type) {
    62  	case Hasher:
    63  		key := msg.Hash()
    64  		hmc := state.hmc
    65  
    66  		node, ok := hmc.hashring.GetNode(key)
    67  		if !ok {
    68  			log.Printf("[ROUTING] Consistent has router failed to derminate routee: %v", key)
    69  			return
    70  		}
    71  		if routee, ok := hmc.routeeMap[node]; ok {
    72  			state.sender.Send(routee, message)
    73  		} else {
    74  			log.Println("[ROUTING] Consistent router failed to resolve node", node)
    75  		}
    76  	default:
    77  		log.Println("[ROUTING] Message must implement router.Hasher", msg)
    78  	}
    79  }
    80  
    81  func (state *consistentHashRouterState) InvokeRouterManagementMessage(msg ManagementMessage, sender *actor.PID) {
    82  }
    83  
    84  func NewConsistentHashPool(size int, opts ...actor.PropsOption) *actor.Props {
    85  	return (&actor.Props{}).
    86  		Configure(actor.WithSpawnFunc(spawner(&consistentHashPoolRouter{PoolRouter{PoolSize: size}}))).
    87  		Configure(opts...)
    88  }
    89  
    90  func NewConsistentHashGroup(routees ...*actor.PID) *actor.Props {
    91  	return (&actor.Props{}).Configure(actor.WithSpawnFunc(spawner(&consistentHashGroupRouter{GroupRouter{Routees: actor.NewPIDSet(routees...)}})))
    92  }
    93  
    94  func (config *consistentHashPoolRouter) CreateRouterState() State {
    95  	return &consistentHashRouterState{}
    96  }
    97  
    98  func (config *consistentHashGroupRouter) CreateRouterState() State {
    99  	return &consistentHashRouterState{}
   100  }