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 }