1
+ // -----------------------------------------------------------------------
2
+ // <copyright file="WrappedShardBufferedMessageSpec.cs" company="Akka.NET Project">
3
+ // Copyright (C) 2009-2025 Lightbend Inc. <http://www.lightbend.com>
4
+ // Copyright (C) 2013-2025 .NET Foundation <https://github.com/akkadotnet/akka.net>
5
+ // </copyright>
6
+ // -----------------------------------------------------------------------
7
+
8
+ using System . Collections . Immutable ;
9
+ using System . Threading . Tasks ;
10
+ using Akka . Actor ;
11
+ using Akka . Cluster . Sharding . Internal ;
12
+ using Akka . Cluster . Tools . Singleton ;
13
+ using Akka . Configuration ;
14
+ using Akka . Event ;
15
+ using Akka . TestKit ;
16
+ using FluentAssertions ;
17
+ using Xunit ;
18
+ using Xunit . Abstractions ;
19
+
20
+ namespace Akka . Cluster . Sharding . Tests ;
21
+
22
+ public class WrappedShardBufferedMessageSpec : AkkaSpec
23
+ {
24
+ #region Custom Classes
25
+
26
+ private sealed class MyEnvelope : IWrappedMessage
27
+ {
28
+ public MyEnvelope ( object message )
29
+ {
30
+ Message = message ;
31
+ }
32
+
33
+ public object Message { get ; }
34
+ }
35
+
36
+ private sealed class BufferMessageAdapter : IShardingBufferMessageAdapter
37
+ {
38
+ public object Apply ( object message , IActorContext context )
39
+ => new MyEnvelope ( message ) ;
40
+ }
41
+
42
+ private class EchoActor : UntypedActor
43
+ {
44
+ private readonly ILoggingAdapter _log = Context . GetLogger ( ) ;
45
+ protected override void OnReceive ( object message )
46
+ {
47
+ _log . Info ( $ ">>>> Received { message } ") ;
48
+ Sender . Tell ( message ) ;
49
+ }
50
+ }
51
+
52
+ #endregion
53
+ private sealed class FakeRememberEntitiesProvider : IRememberEntitiesProvider
54
+ {
55
+ private readonly IActorRef _probe ;
56
+
57
+ public FakeRememberEntitiesProvider ( IActorRef probe )
58
+ {
59
+ _probe = probe ;
60
+ }
61
+
62
+ public Props CoordinatorStoreProps ( ) => FakeCoordinatorStoreActor . Props ( ) ;
63
+
64
+ public Props ShardStoreProps ( string shardId ) => FakeShardStoreActor . Props ( shardId , _probe ) ;
65
+ }
66
+
67
+ private class ShardStoreCreated
68
+ {
69
+ public ShardStoreCreated ( IActorRef store , string shardId )
70
+ {
71
+ Store = store ;
72
+ ShardId = shardId ;
73
+ }
74
+
75
+ public IActorRef Store { get ; }
76
+ public string ShardId { get ; }
77
+ }
78
+
79
+ private class CoordinatorStoreCreated
80
+ {
81
+ public CoordinatorStoreCreated ( IActorRef store )
82
+ {
83
+ Store = store ;
84
+ }
85
+
86
+ public IActorRef Store { get ; }
87
+ }
88
+
89
+ private class FakeShardStoreActor : ActorBase
90
+ {
91
+ public static Props Props ( string shardId , IActorRef probe ) => Actor . Props . Create ( ( ) => new FakeShardStoreActor ( shardId , probe ) ) ;
92
+
93
+ private readonly string _shardId ;
94
+ private readonly IActorRef _probe ;
95
+
96
+ private FakeShardStoreActor ( string shardId , IActorRef probe )
97
+ {
98
+ _shardId = shardId ;
99
+ _probe = probe ;
100
+ Context . System . EventStream . Publish ( new ShardStoreCreated ( Self , shardId ) ) ;
101
+ }
102
+
103
+ protected override bool Receive ( object message )
104
+ {
105
+ switch ( message )
106
+ {
107
+ case RememberEntitiesShardStore . GetEntities :
108
+ Sender . Tell ( new RememberEntitiesShardStore . RememberedEntities ( ImmutableHashSet < string > . Empty ) ) ;
109
+ return true ;
110
+ case RememberEntitiesShardStore . Update m :
111
+ _probe . Tell ( new RememberEntitiesShardStore . UpdateDone ( m . Started , m . Stopped ) ) ;
112
+ return true ;
113
+ }
114
+ return false ;
115
+ }
116
+ }
117
+
118
+ private class FakeCoordinatorStoreActor : ActorBase
119
+ {
120
+ public static Props Props ( ) => Actor . Props . Create ( ( ) => new FakeCoordinatorStoreActor ( ) ) ;
121
+
122
+ private FakeCoordinatorStoreActor ( )
123
+ {
124
+ Context . System . EventStream . Publish ( new CoordinatorStoreCreated ( Context . Self ) ) ;
125
+ }
126
+
127
+ protected override bool Receive ( object message )
128
+ {
129
+ switch ( message )
130
+ {
131
+ case RememberEntitiesCoordinatorStore . GetShards _:
132
+ Sender . Tell ( new RememberEntitiesCoordinatorStore . RememberedShards ( ImmutableHashSet < string > . Empty ) ) ;
133
+ return true ;
134
+ case RememberEntitiesCoordinatorStore . AddShard m :
135
+ Sender . Tell ( new RememberEntitiesCoordinatorStore . UpdateDone ( m . ShardId ) ) ;
136
+ return true ;
137
+ }
138
+ return false ;
139
+ }
140
+ }
141
+
142
+ private static Config GetConfig ( )
143
+ {
144
+ return ConfigurationFactory . ParseString ( @"
145
+ akka.loglevel=DEBUG
146
+ akka.actor.provider = cluster
147
+ akka.remote.dot-netty.tcp.port = 0
148
+ akka.cluster.sharding.state-store-mode = ddata
149
+ akka.cluster.sharding.remember-entities = on
150
+
151
+ # no leaks between test runs thank you
152
+ akka.cluster.sharding.distributed-data.durable.keys = []
153
+ akka.cluster.sharding.verbose-debug-logging = on
154
+ akka.cluster.sharding.fail-on-invalid-entity-state-transition = on" )
155
+
156
+ . WithFallback ( Sharding . ClusterSharding . DefaultConfig ( ) )
157
+ . WithFallback ( DistributedData . DistributedData . DefaultConfig ( ) )
158
+ . WithFallback ( ClusterSingleton . DefaultConfig ( ) ) ;
159
+ }
160
+
161
+ private readonly IActorRef _shard ;
162
+ private IActorRef _store ;
163
+
164
+ public WrappedShardBufferedMessageSpec ( ITestOutputHelper output ) : base ( GetConfig ( ) , output )
165
+ {
166
+ Sys . EventStream . Subscribe ( TestActor , typeof ( ShardStoreCreated ) ) ;
167
+ Sys . EventStream . Subscribe ( TestActor , typeof ( CoordinatorStoreCreated ) ) ;
168
+
169
+ _shard = ChildActorOf ( Shard . Props (
170
+ typeName : "test" ,
171
+ shardId : "test" ,
172
+ entityProps : _ => Props . Create ( ( ) => new EchoActor ( ) ) ,
173
+ settings : ClusterShardingSettings . Create ( Sys ) ,
174
+ extractor : new ExtractorAdapter ( HashCodeMessageExtractor . Create ( 10 , m => m . ToString ( ) ) ) ,
175
+ handOffStopMessage : PoisonPill . Instance ,
176
+ rememberEntitiesProvider : new FakeRememberEntitiesProvider ( TestActor ) ,
177
+ bufferMessageAdapter : new BufferMessageAdapter ( ) ) ) ;
178
+ }
179
+
180
+ private async Task < RememberEntitiesShardStore . UpdateDone > ExpectShardStartup ( )
181
+ {
182
+ var createdEvent = await ExpectMsgAsync < ShardStoreCreated > ( ) ;
183
+ createdEvent . ShardId . Should ( ) . Be ( "test" ) ;
184
+
185
+ _store = createdEvent . Store ;
186
+
187
+ await ExpectMsgAsync < ShardInitialized > ( ) ;
188
+
189
+ _shard . Tell ( new ShardRegion . StartEntity ( "hit" ) ) ;
190
+
191
+ return await ExpectMsgAsync < RememberEntitiesShardStore . UpdateDone > ( ) ;
192
+ }
193
+
194
+ [ Fact ( DisplayName = "Message wrapped in ShardingEnvelope, buffered by Shard, must arrive in entity actor" ) ]
195
+ public async Task WrappedMessageDelivery ( )
196
+ {
197
+ var continueMessage = await ExpectShardStartup ( ) ;
198
+
199
+ // this message should be buffered
200
+ _shard . Tell ( new ShardingEnvelope ( "hit" , "hit" ) ) ;
201
+ await Task . Yield ( ) ;
202
+
203
+ // Tell shard to continue processing
204
+ _shard . Tell ( continueMessage ) ;
205
+
206
+ await ExpectMsgAsync ( "hit" ) ;
207
+ }
208
+ }
0 commit comments