github.com/asynkron/protoactor-go@v0.0.0-20240308120642-ef91a6abee75/cluster/cluster_test_tool/cluster_fixture.go (about) 1 package cluster_test_tool 2 3 import ( 4 "context" 5 "log/slog" 6 "time" 7 8 "github.com/asynkron/protoactor-go/actor" 9 "github.com/asynkron/protoactor-go/cluster" 10 "github.com/asynkron/protoactor-go/cluster/clusterproviders/test" 11 "github.com/asynkron/protoactor-go/cluster/identitylookup/disthash" 12 "github.com/asynkron/protoactor-go/remote" 13 "github.com/google/uuid" 14 "golang.org/x/sync/errgroup" 15 ) 16 17 type ClusterFixture interface { 18 GetMembers() []*cluster.Cluster 19 GetClusterSize() int 20 SpawnNode() *cluster.Cluster 21 RemoveNode(node *cluster.Cluster, graceful bool) 22 ShutDown() 23 } 24 25 type ClusterFixtureConfig struct { 26 GetClusterKinds func() []*cluster.Kind 27 GetClusterProvider func() cluster.ClusterProvider 28 Configure func(*cluster.Config) *cluster.Config 29 GetIdentityLookup func(clusterName string) cluster.IdentityLookup 30 OnDeposing func() 31 } 32 33 type ClusterFixtureOption func(*ClusterFixtureConfig) 34 35 // WithGetClusterKinds sets the cluster kinds for the cluster fixture 36 func WithGetClusterKinds(getKinds func() []*cluster.Kind) ClusterFixtureOption { 37 return func(c *ClusterFixtureConfig) { 38 c.GetClusterKinds = getKinds 39 } 40 } 41 42 // WithClusterConfigure sets the cluster configure function for the cluster fixture 43 func WithClusterConfigure(configure func(*cluster.Config) *cluster.Config) ClusterFixtureOption { 44 return func(c *ClusterFixtureConfig) { 45 c.Configure = configure 46 } 47 } 48 49 // WithGetClusterProvider sets the cluster provider for the cluster fixture 50 func WithGetClusterProvider(getProvider func() cluster.ClusterProvider) ClusterFixtureOption { 51 return func(c *ClusterFixtureConfig) { 52 c.GetClusterProvider = getProvider 53 } 54 } 55 56 // WithGetIdentityLookup sets the identity lookup function for the cluster fixture 57 func WithGetIdentityLookup(identityLookup func(clusterName string) cluster.IdentityLookup) ClusterFixtureOption { 58 return func(c *ClusterFixtureConfig) { 59 c.GetIdentityLookup = identityLookup 60 } 61 } 62 63 // WithOnDeposing sets the on deposing function for the cluster fixture 64 func WithOnDeposing(onDeposing func()) ClusterFixtureOption { 65 return func(c *ClusterFixtureConfig) { 66 c.OnDeposing = onDeposing 67 } 68 } 69 70 const InvalidIdentity string = "invalid" 71 72 type BaseClusterFixture struct { 73 clusterName string 74 clusterSize int 75 config *ClusterFixtureConfig 76 members []*cluster.Cluster 77 } 78 79 func NewBaseClusterFixture(clusterSize int, opts ...ClusterFixtureOption) *BaseClusterFixture { 80 config := &ClusterFixtureConfig{ 81 GetClusterKinds: func() []*cluster.Kind { return make([]*cluster.Kind, 0) }, 82 GetClusterProvider: func() cluster.ClusterProvider { return test.NewTestProvider(test.NewInMemAgent()) }, 83 Configure: func(c *cluster.Config) *cluster.Config { return c }, 84 GetIdentityLookup: func(clusterName string) cluster.IdentityLookup { return disthash.New() }, 85 OnDeposing: func() {}, 86 } 87 for _, opt := range opts { 88 opt(config) 89 } 90 91 fixTure := &BaseClusterFixture{ 92 clusterSize: clusterSize, 93 clusterName: "test-cluster-" + uuid.NewString()[0:6], 94 config: config, 95 members: make([]*cluster.Cluster, 0), 96 } 97 return fixTure 98 } 99 100 // Initialize initializes the cluster fixture 101 func (b *BaseClusterFixture) Initialize() { 102 nodes := b.spawnClusterNodes() 103 b.members = append(b.members, nodes...) 104 } 105 106 func (b *BaseClusterFixture) GetMembers() []*cluster.Cluster { 107 return b.members 108 } 109 110 func (b *BaseClusterFixture) GetClusterSize() int { 111 return b.clusterSize 112 } 113 114 func (b *BaseClusterFixture) SpawnNode() *cluster.Cluster { 115 node := b.spawnClusterMember() 116 b.members = append(b.members, node) 117 return node 118 } 119 120 func (b *BaseClusterFixture) RemoveNode(node *cluster.Cluster, graceful bool) { 121 has := false 122 for i, member := range b.members { 123 if member == node { 124 has = true 125 b.members = append(b.members[:i], b.members[i+1:]...) 126 member.Shutdown(graceful) 127 break 128 } 129 } 130 if !has { 131 slog.Default().Error("node not found", slog.Any("node", node)) 132 } 133 } 134 135 func (b *BaseClusterFixture) ShutDown() { 136 b.config.OnDeposing() 137 b.waitForMembersToShutdown() 138 b.members = b.members[:0] 139 } 140 141 // spawnClusterNodes spawns a number of cluster nodes 142 func (b *BaseClusterFixture) spawnClusterNodes() []*cluster.Cluster { 143 nodes := make([]*cluster.Cluster, 0, b.clusterSize) 144 for i := 0; i < b.clusterSize; i++ { 145 nodes = append(nodes, b.spawnClusterMember()) 146 } 147 148 bgCtx := context.Background() 149 timeoutCtx, cancel := context.WithTimeout(bgCtx, time.Second*10) 150 defer cancel() 151 group := new(errgroup.Group) 152 for _, node := range nodes { 153 tmpNode := node 154 group.Go(func() error { 155 done := make(chan struct{}) 156 go func() { 157 tmpNode.MemberList.TopologyConsensus(timeoutCtx) 158 close(done) 159 }() 160 161 select { 162 case <-timeoutCtx.Done(): 163 return timeoutCtx.Err() 164 case <-done: 165 return nil 166 } 167 }) 168 } 169 err := group.Wait() 170 if err != nil { 171 panic("Failed to reach consensus") 172 } 173 174 return nodes 175 } 176 177 // spawnClusterMember spawns a cluster members 178 func (b *BaseClusterFixture) spawnClusterMember() *cluster.Cluster { 179 config := cluster.Configure(b.clusterName, b.config.GetClusterProvider(), b.config.GetIdentityLookup(b.clusterName), 180 remote.Configure("localhost", 0), 181 cluster.WithKinds(b.config.GetClusterKinds()...), 182 ) 183 config = b.config.Configure(config) 184 185 system := actor.NewActorSystem() 186 187 c := cluster.New(system, config) 188 c.StartMember() 189 return c 190 } 191 192 // waitForMembersToShutdown waits for the members to shutdown 193 func (b *BaseClusterFixture) waitForMembersToShutdown() { 194 for _, member := range b.members { 195 slog.Default().Info("Preparing shutdown for cluster member", slog.String("member", member.ActorSystem.ID)) 196 } 197 198 group := new(errgroup.Group) 199 timeoutCtx, cancel := context.WithTimeout(context.Background(), time.Second*1000) 200 defer cancel() 201 202 for _, member := range b.members { 203 member := member 204 group.Go(func() error { 205 done := make(chan struct{}) 206 go func() { 207 slog.Default().Info("Shutting down cluster member", slog.String("member", member.ActorSystem.ID)) 208 member.Shutdown(true) 209 close(done) 210 }() 211 212 select { 213 case <-timeoutCtx.Done(): 214 return timeoutCtx.Err() 215 case <-done: 216 return nil 217 } 218 }) 219 } 220 err := group.Wait() 221 if err != nil { 222 panic(err) 223 } 224 }