2121import static java .nio .charset .StandardCharsets .UTF_8 ;
2222import static java .util .concurrent .CompletableFuture .completedFuture ;
2323import static java .util .concurrent .CompletableFuture .supplyAsync ;
24+ import static java .util .stream .Collectors .groupingBy ;
2425import static java .util .stream .Collectors .toList ;
2526import static java .util .stream .IntStream .range ;
2627import static org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .base .Preconditions .checkNotNull ;
3031import static org .junit .Assert .assertEquals ;
3132import static org .junit .Assert .assertTrue ;
3233import static org .mockito .ArgumentMatchers .any ;
34+ import static org .mockito .Mockito .atLeast ;
3335import static org .mockito .Mockito .atLeastOnce ;
3436import static org .mockito .Mockito .times ;
3537import static org .mockito .Mockito .verify ;
3840
3941import java .util .Arrays ;
4042import java .util .HashSet ;
43+ import java .util .List ;
4144import java .util .Map ;
4245import java .util .Set ;
4346import java .util .stream .Collectors ;
6063import org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .collect .ImmutableMap ;
6164import org .apache .beam .vendor .guava .v32_1_2_jre .com .google .common .collect .Streams ;
6265import org .apache .commons .lang3 .RandomUtils ;
66+ import org .checkerframework .checker .nullness .qual .Nullable ;
6367import org .joda .time .Duration ;
6468import org .junit .Before ;
6569import org .junit .Rule ;
@@ -263,10 +267,16 @@ public void testWriteBatchesWithTimeout() {
263267
264268 p .run ().waitUntilFinish ();
265269
266- SendMessageBatchRequestEntry [] entries = entries (range (0 , 5 ));
267- // due to added delay, batches are timed out on arrival of every 3rd msg
268- verify (sqs ).sendMessageBatch (request ("queue" , entries [0 ], entries [1 ], entries [2 ]));
269- verify (sqs ).sendMessageBatch (request ("queue" , entries [3 ], entries [4 ]));
270+ // Nominally batches time out on arrival of every 3rd message ([0,1,2], [3,4]). The exact
271+ // grouping depends on wall clock time and is unreliable on loaded machines (#38946), so
272+ // verify timing-independent invariants instead: expired batches are flushed on append, so
273+ // with >= 100ms between messages no batch can exceed 3 entries before exceeding the 150ms
274+ // timeout, forcing at least 2 batches for 5 messages.
275+ Map <String , List <SendMessageBatchRequest >> requests = captureBatchRequests (atLeast (2 ));
276+ assertThat (requests .keySet ()).containsExactly ("queue" );
277+ assertMessageBodies (requests .get ("queue" ), range (0 , 5 ));
278+ assertThat (requests .get ("queue" ))
279+ .allSatisfy (req -> assertThat (req .entries ().size ()).isLessThanOrEqualTo (3 ));
270280 }
271281
272282 @ Test
@@ -285,11 +295,16 @@ public void testWriteBatchesWithStrictTimeout() {
285295
286296 p .run ().waitUntilFinish ();
287297
288- SendMessageBatchRequestEntry [] entries = entries (range (0 , 5 ));
289- // using strict timeouts batches, batches are timed out by a separate thread
290- verify (sqs ).sendMessageBatch (request ("queue" , entries [0 ], entries [1 ]));
291- verify (sqs ).sendMessageBatch (request ("queue" , entries [2 ], entries [3 ]));
292- verify (sqs ).sendMessageBatch (request ("queue" , entries [4 ]));
298+ // Nominally the separate timeout thread flushes [0,1], [2,3], [4]. The exact grouping
299+ // depends on wall clock time and is unreliable on loaded machines (#38946), so verify
300+ // timing-independent invariants instead. Expired batches are also flushed on append
301+ // (independently of the timeout thread), so no batch can exceed 3 entries and 5 messages
302+ // require at least 2 batches.
303+ Map <String , List <SendMessageBatchRequest >> requests = captureBatchRequests (atLeast (2 ));
304+ assertThat (requests .keySet ()).containsExactly ("queue" );
305+ assertMessageBodies (requests .get ("queue" ), range (0 , 5 ));
306+ assertThat (requests .get ("queue" ))
307+ .allSatisfy (req -> assertThat (req .entries ().size ()).isLessThanOrEqualTo (3 ));
293308 }
294309
295310 @ Test
@@ -337,11 +352,18 @@ public void testWriteBatchesToDynamicWithTimeout() {
337352
338353 p .run ().waitUntilFinish ();
339354
340- SendMessageBatchRequestEntry [] entries = entries (range (0 , 5 ));
341- // due to added delay, dynamic batches are timed out on arrival of every 2nd msg (per batch)
342- verify (sqs ).sendMessageBatch (request ("even" , entries [0 ], entries [2 ]));
343- verify (sqs ).sendMessageBatch (request ("uneven" , entries [1 ], entries [3 ]));
344- verify (sqs ).sendMessageBatch (request ("even" , entries [4 ]));
355+ // Nominally dynamic batches are timed out on arrival of every 2nd message per batch
356+ // ([0,2], [1,3], [4]). The exact grouping depends on wall clock time and is unreliable on
357+ // loaded machines (#38946), so verify timing-independent invariants instead: per queue,
358+ // consecutive messages are >= 200ms apart, so expired batches are flushed on append once
359+ // reaching 2 entries, requiring at least 3 batches overall.
360+ Map <String , List <SendMessageBatchRequest >> requests = captureBatchRequests (atLeast (3 ));
361+ assertThat (requests .keySet ()).containsExactlyInAnyOrder ("even" , "uneven" );
362+ assertMessageBodies (requests .get ("even" ), range (0 , 5 ).filter (i -> i % 2 == 0 ));
363+ assertMessageBodies (requests .get ("uneven" ), range (0 , 5 ).filter (i -> i % 2 == 1 ));
364+ requests .values ().stream ()
365+ .flatMap (List ::stream )
366+ .forEach (req -> assertThat (req .entries ().size ()).isLessThanOrEqualTo (2 ));
345367 }
346368
347369 @ Test
@@ -360,14 +382,19 @@ public void testWriteBatchesToDynamicWithStrictTimeout() {
360382
361383 p .run ().waitUntilFinish ();
362384
363- SendMessageBatchRequestEntry [] entries = entries (range (0 , 5 ));
364- // using strict timeouts batches, batches are timed out by a separate thread before any 2nd
365- // entry
366- verify (sqs ).sendMessageBatch (request ("even" , entries [0 ]));
367- verify (sqs ).sendMessageBatch (request ("uneven" , entries [1 ]));
368- verify (sqs ).sendMessageBatch (request ("even" , entries [2 ]));
369- verify (sqs ).sendMessageBatch (request ("uneven" , entries [3 ]));
370- verify (sqs ).sendMessageBatch (request ("even" , entries [4 ]));
385+ // Nominally the separate timeout thread flushes every batch before a 2nd entry arrives
386+ // (5 singleton batches). The exact grouping depends on wall clock time and is unreliable
387+ // on loaded machines (#38946), so verify timing-independent invariants instead. Expired
388+ // batches are also flushed on append (independently of the timeout thread) and per queue
389+ // consecutive messages are >= 200ms apart, so no batch can exceed 2 entries and at least
390+ // 3 batches are required overall.
391+ Map <String , List <SendMessageBatchRequest >> requests = captureBatchRequests (atLeast (3 ));
392+ assertThat (requests .keySet ()).containsExactlyInAnyOrder ("even" , "uneven" );
393+ assertMessageBodies (requests .get ("even" ), range (0 , 5 ).filter (i -> i % 2 == 0 ));
394+ assertMessageBodies (requests .get ("uneven" ), range (0 , 5 ).filter (i -> i % 2 == 1 ));
395+ requests .values ().stream ()
396+ .flatMap (List ::stream )
397+ .forEach (req -> assertThat (req .entries ().size ()).isLessThanOrEqualTo (2 ));
371398 }
372399
373400 @ Test
@@ -406,6 +433,28 @@ private SendMessageBatchRequest anyRequest() {
406433 return any ();
407434 }
408435
436+ /** Captures all batch requests, verifying the given mode, and groups them by queue url. */
437+ private Map <String , List <SendMessageBatchRequest >> captureBatchRequests (
438+ org .mockito .verification .VerificationMode mode ) {
439+ ArgumentCaptor <SendMessageBatchRequest > captor =
440+ ArgumentCaptor .forClass (SendMessageBatchRequest .class );
441+ verify (sqs , mode ).sendMessageBatch (captor .capture ());
442+ return captor .getAllValues ().stream ().collect (groupingBy (SendMessageBatchRequest ::queueUrl ));
443+ }
444+
445+ /** Asserts that the requests contain exactly the expected message bodies, each exactly once. */
446+ private void assertMessageBodies (
447+ @ Nullable List <SendMessageBatchRequest > requests , IntStream expectedMsgs ) {
448+ assertThat (requests ).isNotNull ();
449+ assertThat (
450+ requests .stream ()
451+ .flatMap (req -> req .entries ().stream ())
452+ .map (SendMessageBatchRequestEntry ::messageBody )
453+ .collect (toList ()))
454+ .containsExactlyInAnyOrderElementsOf (
455+ expectedMsgs .mapToObj (Integer ::toString ).collect (toList ()));
456+ }
457+
409458 private SendMessageBatchRequest request (String queue , SendMessageBatchRequestEntry ... entries ) {
410459 return SendMessageBatchRequest .builder ()
411460 .queueUrl (queue )
0 commit comments