Fix multi-stream

etspaceman · Jan 23, 2025 · 5e6695b · 5e6695b
1 parent 252aa16
commit 5e6695b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/kcl-localstack/src/main/scala/kinesis4cats/kcl/fs2/localstack/LocalstackKCLConsumerFS2.scala b/kcl-localstack/src/main/scala/kinesis4cats/kcl/fs2/localstack/LocalstackKCLConsumerFS2.scala
@@ -70,8 +70,11 @@ object LocalstackKCLConsumerFS2 {
           _.configureLeaseManagementConfig(
             LocalstackKCLConsumer.configureTopLevelLeaseManagementConfig
           )
-            .configureLeaseManagementConfig(
-              LocalstackKCLConsumer.configureLeaseManagementFactory
+            .configureLeaseManagementConfig(x =>
+              LocalstackKCLConsumer.configureLeaseManagementFactory(
+                x,
+                streamTracker.isMultiStream()
+              )
             )
             .configureCoordinatorConfig(_.parentShardPollIntervalMillis(1000L))
             .configureRetrievalConfig(

diff --git a/kcl-localstack/src/main/scala/kinesis4cats/kcl/localstack/LocalstackKCLConsumer.scala b/kcl-localstack/src/main/scala/kinesis4cats/kcl/localstack/LocalstackKCLConsumer.scala
@@ -45,7 +45,8 @@ object LocalstackKCLConsumer {
       .failoverTimeMillis(1000L)
 
   private[kcl] def configureLeaseManagementFactory(
-      defaultLeaseManagement: LeaseManagementConfig
+      defaultLeaseManagement: LeaseManagementConfig,
+      isMultiStream: Boolean
   ): LeaseManagementConfig =
     defaultLeaseManagement.leaseManagementFactory(
       new DynamoDBLeaseManagementFactory(
@@ -79,7 +80,7 @@ object LocalstackKCLConsumer {
         defaultLeaseManagement.tags(),
         new DynamoDBLeaseSerializer(),
         defaultLeaseManagement.customShardDetectorProvider(),
-        false,
+        isMultiStream,
         LeaseCleanupConfig
           .builder()
           .completedLeaseCleanupIntervalMillis(500L)
@@ -162,7 +163,9 @@ object LocalstackKCLConsumer {
             .configureLeaseManagementConfig(
               configureTopLevelLeaseManagementConfig
             )
-            .configureLeaseManagementConfig(configureLeaseManagementFactory)
+            .configureLeaseManagementConfig(x =>
+              configureLeaseManagementFactory(x, streamTracker.isMultiStream())
+            )
             .configureCoordinatorConfig(_.parentShardPollIntervalMillis(1000L))
             .configureRetrievalConfig(
               _.retrievalSpecificConfig(retrievalConfig)