Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve Test Reliability by Replacing ZIO.sleep with Synchronization Primitives in ConsumerSpec #1401

Merged
merged 3 commits into from
Nov 23, 2024
Merged

Improve Test Reliability by Replacing ZIO.sleep with Synchronization Primitives in ConsumerSpec #1401

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
87e0a5e
Use Promises to improve flaky tests
AdrielC Nov 22, 2024
18af699
fmt
AdrielC Nov 22, 2024
730ebc1
Merge branch 'master' into master
erikvanoosten Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 46 additions & 17 deletions zio-kafka-test/src/test/scala/zio/kafka/consumer/ConsumerSpec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,7 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {
}.tap { case (_, idx) => ZIO.logDebug(s"Consumed $idx") }
}
.runDrain
.tap(_ => ZIO.logDebug("Stream completed"))
.zipLeft(ZIO.logDebug("Stream completed"))
.provideSomeLayer[Kafka](
consumer(client, Some(group))
)
Expand Down Expand Up @@ -530,18 +530,30 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {

// Consume messages
subscription = Subscription.topics(topic)
assignedPartitionsRef <- Ref.make(Set.empty[Int]) // Set of partition numbers
// Create a Promise to signal when consumer1 has processed half the partitions
consumer1Ready <- Promise.make[Nothing, Unit]
consumer1 <- Consumer
.partitionedStream(subscription, Serde.string, Serde.string)
.flatMapPar(nrPartitions) { case (tp, partition) =>
ZStream
.fromZIO(partition.runDrain)
.fromZIO(
consumer1Ready
.succeed(())
.whenZIO(
assignedPartitionsRef
.updateAndGet(_ + tp.partition())
.map(_.size >= (nrPartitions / 2))
) *>
partition.runDrain
)
.as(tp)
}
.take(nrPartitions.toLong / 2)
.runDrain
.provideSomeLayer[Kafka](consumer(client1, Some(group)))
.fork
_ <- Live.live(ZIO.sleep(5.seconds))
_ <- consumer1Ready.await
consumer2 <- Consumer
.partitionedStream(subscription, Serde.string, Serde.string)
.take(nrPartitions.toLong / 2)
Expand Down Expand Up @@ -574,11 +586,22 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {

// Consume messages
subscription = Subscription.topics(topic)
consumer1Ready <- Promise.make[Nothing, Unit]
assignedPartitionsRef <- Ref.make(Set.empty[Int]) // Set of partition numbers
consumer1 <- Consumer
.partitionedStream(subscription, Serde.string, Serde.string)
.flatMapPar(nrPartitions) { case (tp, partition) =>
ZStream
.fromZIO(partition.runDrain)
.fromZIO(
consumer1Ready
.succeed(())
.whenZIO(
assignedPartitionsRef
.updateAndGet(_ + tp.partition())
.map(_.size >= (nrPartitions / 2))
) *>
partition.runDrain
)
.as(tp)
}
.take(nrPartitions.toLong / 2)
Expand All @@ -592,15 +615,14 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {
.collect { case rebalance: DiagnosticEvent.Rebalance => rebalance }
.runCollect
.fork
_ <- ZIO.sleep(5.seconds)
_ <- consumer1Ready.await
consumer2 <- Consumer
.partitionedStream(subscription, Serde.string, Serde.string)
.take(nrPartitions.toLong / 2)
.runDrain
.provideSomeLayer[Kafka](consumer(client2, Some(group)))
.fork
_ <- consumer1.join
_ <- consumer1.join
_ <- consumer2.join
} yield diagnosticStream.join
}
Expand Down Expand Up @@ -1481,6 +1503,7 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {
"it's possible to start a new consumption session from a Consumer that had a consumption session stopped previously"
) {
val numberOfMessages: Int = 100000
val messagesToConsumeBeforeStop = 1000 // Adjust this threshold as needed
AdrielC marked this conversation as resolved.
Show resolved Hide resolved
val kvs: Iterable[(String, String)] = Iterable.tabulate(numberOfMessages)(i => (s"key-$i", s"msg-$i"))

def test(diagnostics: Diagnostics): ZIO[Producer & Scope & Kafka, Throwable, TestResult] =
Expand All @@ -1490,22 +1513,28 @@ object ConsumerSpec extends ZIOSpecDefaultSlf4j with KafkaRandom {
settings <- consumerSettings(clientId = clientId)
consumer <- Consumer.make(settings, diagnostics = diagnostics)
_ <- produceMany(topic, kvs)
// Create a Ref to track messages consumed and a Promise to signal when to stop consumption
messagesConsumedRef <- Ref.make(0)
stopPromise <- Promise.make[Nothing, Unit]
// Starting a consumption session to start the Runloop.
fiber <-
consumer
.plainStream(Subscription.manual(topic -> 0), Serde.string, Serde.string)
.tap(_ => ZIO.sleep(1.millisecond)) // sleep to avoid consuming all messages in under 200 millis
.take(numberOfMessages.toLong)
.runCount
.forkScoped
_ <- ZIO.sleep(200.millis)
fiber <- consumer
.plainStream(Subscription.manual(topic -> 0), Serde.string, Serde.string)
.mapZIO { _ =>
messagesConsumedRef.updateAndGet(_ + 1).flatMap { count =>
if (count >= messagesToConsumeBeforeStop) stopPromise.succeed(()).as(1L)
else ZIO.succeed(1L)
}
}
.take(numberOfMessages.toLong)
.runSum
.forkScoped

// Wait for the consumption to reach the desired threshold
_ <- stopPromise.await
_ <- consumer.stopConsumption
consumed0 <- fiber.join
_ <- ZIO.logDebug(s"consumed0: $consumed0")

_ <- ZIO.logDebug("About to sleep 5 seconds")
_ <- ZIO.sleep(5.seconds)
_ <- ZIO.logDebug("Slept 5 seconds")
consumed1 <- consumer
.plainStream(Subscription.manual(topic -> 0), Serde.string, Serde.string)
.take(numberOfMessages.toLong)
Expand Down