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  }