@@ -34,7 +34,7 @@ import com.snowplowanalytics.snowplow.pubsub.GcpUserAgent
34
34
import com .snowplowanalytics .snowplow .sources .SourceAndAck
35
35
import com .snowplowanalytics .snowplow .sources .internal .{Checkpointer , LowLevelEvents , LowLevelSource }
36
36
37
- import scala .concurrent .duration .{Duration , DurationDouble , FiniteDuration }
37
+ import scala .concurrent .duration .{DurationDouble , FiniteDuration }
38
38
import scala .jdk .CollectionConverters ._
39
39
40
40
import java .util .concurrent .{ExecutorService , Executors , LinkedBlockingQueue }
@@ -92,7 +92,7 @@ object PubsubSourceV2 {
92
92
(hotswap, _) <- Stream .resource(Hotswap (resource))
93
93
fs2Queue <- Stream .eval(Queue .synchronous[F , SubscriberAction ])
94
94
_ <- extendDeadlines(config, stub, refStates, channelAffinity).spawn
95
- _ <- Stream .eval(queueToQueue(config, jQueue, fs2Queue, stub, channelAffinity )).repeat.spawn
95
+ _ <- Stream .eval(queueToQueue(config, jQueue, fs2Queue, channelAffinity, hotswap )).repeat.spawn
96
96
lle <- Stream
97
97
.fromQueueUnterminated(fs2Queue)
98
98
.through(toLowLevelEvents(config, refStates, hotswap, resource, channelAffinity))
@@ -103,26 +103,18 @@ object PubsubSourceV2 {
103
103
config : PubsubSourceConfigV2 ,
104
104
jQueue : LinkedBlockingQueue [SubscriberAction ],
105
105
fs2Queue : QueueSink [F , SubscriberAction ],
106
- stub : SubscriberStub ,
107
- channelAffinity : Int
106
+ channelAffinity : Int ,
107
+ hotswap : Hotswap [ F , KeepAlive [ F ]]
108
108
): F [Unit ] =
109
109
resolveNextAction(jQueue).flatMap {
110
- case action @ SubscriberAction .ProcessRecords (records, controller, _) =>
111
- val fallback = if (config.modackOnProgressTimeout) {
112
- val ackIds = records.map(_.getAckId)
113
- if (config.cancelOnProgressTimeout)
114
- Logger [F ].debug(s " Cancelling Pubsub channel $channelAffinity for not making progress " ) *>
115
- Sync [F ].delay(controller.cancel()) *> Utils .modAck(config.subscription, stub, ackIds, Duration .Zero , channelAffinity)
116
- else
117
- Logger [F ].debug(s " Nacking on Pubsub channel $channelAffinity for not making progress " ) *>
118
- Sync [F ].delay(controller.request(1 )) *> Utils .modAck(config.subscription, stub, ackIds, Duration .Zero , channelAffinity)
119
- } else {
120
- if (config.cancelOnProgressTimeout)
121
- Logger [F ].debug(s " Cancelling Pubsub channel $channelAffinity for not making progress " ) *>
122
- Sync [F ].delay(controller.cancel()) *> fs2Queue.offer(action)
123
- else
124
- fs2Queue.offer(action)
125
- }
110
+ case action @ SubscriberAction .ProcessRecords (_, _, _) =>
111
+ def fallback : F [Unit ] =
112
+ hotswap.get.use {
113
+ case Some (keepAlive) =>
114
+ Logger [F ].debug(s " Sending keepalie for channel $channelAffinity" ) *>
115
+ keepAlive.keepAlive
116
+ case None => Sync [F ].unit
117
+ } >> fs2Queue.offer(action).timeoutTo(config.progressTimeout, fallback)
126
118
fs2Queue.offer(action).timeoutTo(config.progressTimeout, fallback)
127
119
case action : SubscriberAction .SubscriberError =>
128
120
fs2Queue.offer(action)
@@ -194,8 +186,8 @@ object PubsubSourceV2 {
194
186
private def toLowLevelEvents [F [_]: Async ](
195
187
config : PubsubSourceConfigV2 ,
196
188
refStates : Ref [F , Map [Unique .Token , PubsubBatchState ]],
197
- hotswap : Hotswap [F , Unit ],
198
- toSwap : Resource [F , Unit ],
189
+ hotswap : Hotswap [F , KeepAlive [ F ] ],
190
+ toSwap : Resource [F , KeepAlive [ F ] ],
199
191
channelAffinity : Int
200
192
): Pipe [F , SubscriberAction , LowLevelEvents [Vector [Unique .Token ]]] =
201
193
_.flatMap {
@@ -282,7 +274,7 @@ object PubsubSourceV2 {
282
274
actionQueue : LinkedBlockingQueue [SubscriberAction ],
283
275
channelAffinity : Int ,
284
276
clientId : UUID
285
- ): Resource [F , Unit ] = {
277
+ ): Resource [F , KeepAlive [ F ] ] = {
286
278
287
279
val observer = new ResponseObserver [StreamingPullResponse ] {
288
280
var controller : StreamController = _
@@ -314,7 +306,7 @@ object PubsubSourceV2 {
314
306
val request = StreamingPullRequest .newBuilder
315
307
.setSubscription(config.subscription.show)
316
308
.setStreamAckDeadlineSeconds(config.durationPerAckExtension.toSeconds.toInt)
317
- .setClientId(if (config.consistentClientId) clientId.toString else UUID .randomUUID .toString)
309
+ .setClientId(clientId.toString)
318
310
.setMaxOutstandingMessages(0 )
319
311
.setMaxOutstandingBytes(0 )
320
312
.build
@@ -323,11 +315,18 @@ object PubsubSourceV2 {
323
315
.make(Sync [F ].delay(subStub.streamingPullCallable.splitCall(observer, context))) { stream =>
324
316
Sync [F ].delay(stream.closeSendWithError(Status .CANCELLED .asException))
325
317
}
326
- .evalMap { stream =>
327
- Sync [F ].delay(stream.send(request))
318
+ .map { stream =>
319
+ new KeepAlive [F ] {
320
+ def keepAlive : F [Unit ] =
321
+ Sync [F ].delay(stream.send(request))
322
+ }
328
323
}
329
- .void
324
+ .evalTap(_.keepAlive)
325
+
326
+ }
330
327
328
+ trait KeepAlive [F [_]] {
329
+ def keepAlive : F [Unit ]
331
330
}
332
331
333
332
private def executorResource [F [_]: Sync , E <: ExecutorService ](make : F [E ]): Resource [F , E ] =
0 commit comments