github.com/asynkron/protoactor-go@v0.0.0-20240308120642-ef91a6abee75/cluster/cluster_test_tool/pubsub_cluster_fixture.go (about)

     1  package cluster_test_tool
     2  
     3  import (
     4  	"errors"
     5  	"log/slog"
     6  	"math/rand"
     7  	"strconv"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/asynkron/protoactor-go/actor"
    13  	"github.com/asynkron/protoactor-go/cluster"
    14  	"github.com/stretchr/testify/assert"
    15  	"golang.org/x/net/context"
    16  	"google.golang.org/protobuf/proto"
    17  )
    18  
    19  const (
    20  	PubSubSubscriberKind        = "Subscriber"
    21  	PubSubTimeoutSubscriberKind = "TimeoutSubscriber"
    22  )
    23  
    24  type PubSubClusterFixture struct {
    25  	*BaseClusterFixture
    26  
    27  	useDefaultTopicRegistration bool
    28  	t                           testing.TB
    29  
    30  	Deliveries     []Delivery
    31  	DeliveriesLock *sync.RWMutex
    32  
    33  	subscriberStore cluster.KeyValueStore[*cluster.Subscribers]
    34  }
    35  
    36  func NewPubSubClusterFixture(t testing.TB, clusterSize int, useDefaultTopicRegistration bool, opts ...ClusterFixtureOption) *PubSubClusterFixture {
    37  	lock := &sync.RWMutex{}
    38  	store := NewInMemorySubscriberStore()
    39  	fixture := &PubSubClusterFixture{
    40  		t:                           t,
    41  		useDefaultTopicRegistration: useDefaultTopicRegistration,
    42  		Deliveries:                  []Delivery{},
    43  		DeliveriesLock:              lock,
    44  		subscriberStore:             store,
    45  	}
    46  
    47  	pubSubOpts := []ClusterFixtureOption{
    48  		WithGetClusterKinds(func() []*cluster.Kind {
    49  			kinds := []*cluster.Kind{
    50  				cluster.NewKind(PubSubSubscriberKind, fixture.subscriberProps()),
    51  				cluster.NewKind(PubSubTimeoutSubscriberKind, fixture.timeoutSubscriberProps()),
    52  			}
    53  			if !fixture.useDefaultTopicRegistration {
    54  				kinds = append(kinds, cluster.NewKind(cluster.TopicActorKind, actor.PropsFromProducer(func() actor.Actor {
    55  					return cluster.NewTopicActor(store, slog.Default())
    56  				})))
    57  			}
    58  			return kinds
    59  		}),
    60  		WithClusterConfigure(func(config *cluster.Config) *cluster.Config {
    61  			cluster.WithRequestTimeout(time.Second * 1)(config)
    62  			cluster.WithPubSubSubscriberTimeout(time.Second * 2)(config)
    63  			return config
    64  		}),
    65  	}
    66  	pubSubOpts = append(pubSubOpts, opts...)
    67  
    68  	fixture.BaseClusterFixture = NewBaseInMemoryClusterFixture(clusterSize, pubSubOpts...)
    69  	return fixture
    70  }
    71  
    72  func (p *PubSubClusterFixture) RandomMember() *cluster.Cluster {
    73  	members := p.BaseClusterFixture.GetMembers()
    74  	return members[rand.Intn(len(members))]
    75  }
    76  
    77  // VerifyAllSubscribersGotAllTheData verifies that all subscribers got all the data
    78  func (p *PubSubClusterFixture) VerifyAllSubscribersGotAllTheData(subscriberIds []string, numMessages int) {
    79  	WaitUntil(p.t, func() bool {
    80  		p.DeliveriesLock.RLock()
    81  		defer p.DeliveriesLock.RUnlock()
    82  		return len(p.Deliveries) == numMessages*len(subscriberIds)
    83  	}, "All messages should be delivered ", DefaultWaitTimeout*1000)
    84  
    85  	p.DeliveriesLock.RLock()
    86  	defer p.DeliveriesLock.RUnlock()
    87  
    88  	expected := make([]Delivery, 0, len(subscriberIds))
    89  	for _, subscriberId := range subscriberIds {
    90  		for i := 0; i < numMessages; i++ {
    91  			expected = append(expected, Delivery{
    92  				Identity: subscriberId,
    93  				Data:     i,
    94  			})
    95  		}
    96  	}
    97  	assert.ElementsMatch(p.t, expected, p.Deliveries)
    98  }
    99  
   100  // SubscribeAllTo subscribes all the given subscribers to the given topic
   101  func (p *PubSubClusterFixture) SubscribeAllTo(topic string, subscriberIds []string) {
   102  	for _, subscriberId := range subscriberIds {
   103  		p.SubscribeTo(topic, subscriberId, PubSubSubscriberKind)
   104  	}
   105  }
   106  
   107  // UnSubscribeAllFrom unsubscribes all the given subscribers from the given topic
   108  func (p *PubSubClusterFixture) UnSubscribeAllFrom(topic string, subscriberIds []string) {
   109  	for _, subscriberId := range subscriberIds {
   110  		p.UnSubscribeTo(topic, subscriberId, PubSubSubscriberKind)
   111  	}
   112  }
   113  
   114  // SubscribeTo subscribes the given subscriber to the given topic
   115  func (p *PubSubClusterFixture) SubscribeTo(topic, identity, kind string) {
   116  	c := p.RandomMember()
   117  	res, err := c.SubscribeByClusterIdentity(topic, cluster.NewClusterIdentity(identity, kind), cluster.WithTimeout(time.Second*5))
   118  	assert.NoError(p.t, err, kind+"/"+identity+" should be able to subscribe to topic "+topic)
   119  	assert.NotNil(p.t, res, kind+"/"+identity+" subscribing should not time out on topic "+topic)
   120  }
   121  
   122  // UnSubscribeTo unsubscribes the given subscriber from the given topic
   123  func (p *PubSubClusterFixture) UnSubscribeTo(topic, identity, kind string) {
   124  	c := p.RandomMember()
   125  	res, err := c.UnsubscribeByClusterIdentity(topic, cluster.NewClusterIdentity(identity, kind), cluster.WithTimeout(time.Second*5))
   126  	assert.NoError(p.t, err, kind+"/"+identity+" should be able to unsubscribe from topic "+topic)
   127  	assert.NotNil(p.t, res, kind+"/"+identity+" subscribing should not time out on topic "+topic)
   128  }
   129  
   130  // PublishData publishes the given message to the given topic
   131  func (p *PubSubClusterFixture) PublishData(topic string, data int) (*cluster.PublishResponse, error) {
   132  	c := p.RandomMember()
   133  	return c.Publisher().Publish(context.Background(), topic, &DataPublished{Data: int32(data)}, cluster.WithTimeout(time.Second*5))
   134  }
   135  
   136  // PublishDataBatch publishes the given messages to the given topic
   137  func (p *PubSubClusterFixture) PublishDataBatch(topic string, data []int) (*cluster.PublishResponse, error) {
   138  	batches := make([]proto.Message, 0)
   139  	for _, d := range data {
   140  		batches = append(batches, &DataPublished{Data: int32(d)})
   141  	}
   142  
   143  	c := p.RandomMember()
   144  	return c.Publisher().PublishBatch(context.Background(), topic, &cluster.PubSubBatch{Envelopes: batches}, cluster.WithTimeout(time.Second*5))
   145  }
   146  
   147  // SubscriberIds returns the subscriber ids
   148  func (p *PubSubClusterFixture) SubscriberIds(prefix string, count int) []string {
   149  	ids := make([]string, 0, count)
   150  	for i := 0; i < count; i++ {
   151  		ids = append(ids, prefix+strconv.Itoa(i))
   152  	}
   153  	return ids
   154  }
   155  
   156  // GetSubscribersForTopic returns the subscribers for the given topic
   157  func (p *PubSubClusterFixture) GetSubscribersForTopic(topic string) (*cluster.Subscribers, error) {
   158  	return p.subscriberStore.Get(context.Background(), topic)
   159  }
   160  
   161  // ClearDeliveries clears the deliveries
   162  func (p *PubSubClusterFixture) ClearDeliveries() {
   163  	p.DeliveriesLock.Lock()
   164  	defer p.DeliveriesLock.Unlock()
   165  	p.Deliveries = make([]Delivery, 0)
   166  }
   167  
   168  // subscriberProps returns the props for the subscriber actor
   169  func (p *PubSubClusterFixture) subscriberProps() *actor.Props {
   170  	return actor.PropsFromFunc(func(context actor.Context) {
   171  		if msg, ok := context.Message().(*DataPublished); ok {
   172  			identity := cluster.GetClusterIdentity(context)
   173  
   174  			p.AppendDelivery(Delivery{
   175  				Identity: identity.Identity,
   176  				Data:     int(msg.Data),
   177  			})
   178  			context.Respond(&Response{})
   179  		}
   180  	})
   181  }
   182  
   183  // timeoutSubscriberProps returns the props for the subscriber actor
   184  func (p *PubSubClusterFixture) timeoutSubscriberProps() *actor.Props {
   185  	return actor.PropsFromFunc(func(context actor.Context) {
   186  		if msg, ok := context.Message().(*DataPublished); ok {
   187  			time.Sleep(time.Second * 4) // 4 seconds is longer than the configured subscriber timeout
   188  
   189  			identity := cluster.GetClusterIdentity(context)
   190  			p.AppendDelivery(Delivery{
   191  				Identity: identity.Identity,
   192  				Data:     int(msg.Data),
   193  			})
   194  			context.Respond(&Response{})
   195  		}
   196  	})
   197  }
   198  
   199  // AppendDelivery appends a delivery to the deliveries slice
   200  func (p *PubSubClusterFixture) AppendDelivery(delivery Delivery) {
   201  	p.DeliveriesLock.Lock()
   202  	p.Deliveries = append(p.Deliveries, delivery)
   203  	p.DeliveriesLock.Unlock()
   204  }
   205  
   206  type Delivery struct {
   207  	Identity string
   208  	Data     int
   209  }
   210  
   211  func NewInMemorySubscriberStore() *InMemorySubscribersStore[*cluster.Subscribers] {
   212  	return &InMemorySubscribersStore[*cluster.Subscribers]{
   213  		store: &sync.Map{},
   214  	}
   215  }
   216  
   217  type InMemorySubscribersStore[T any] struct {
   218  	store *sync.Map // map[string]T
   219  }
   220  
   221  func (i *InMemorySubscribersStore[T]) Set(_ context.Context, key string, value T) error {
   222  	i.store.Store(key, value)
   223  	return nil
   224  }
   225  
   226  func (i *InMemorySubscribersStore[T]) Get(_ context.Context, key string) (T, error) {
   227  	var r T
   228  	value, ok := i.store.Load(key)
   229  	if !ok {
   230  		return r, errors.New("not found")
   231  	}
   232  	return value.(T), nil
   233  }
   234  
   235  func (i *InMemorySubscribersStore[T]) Clear(_ context.Context, key string) error {
   236  	i.store.Delete(key)
   237  	return nil
   238  }