fixed kafka stress tests

Author: aIbrahiim

it/kafka/src/test/java/org/apache/beam/it/kafka/KafkaIOST.java

@@ -155,7 +155,7 @@ public void setup() {
                   Configuration.class),
               "large",
               Configuration.fromJsonString(
-                  "{\"rowsPerSecond\":50000,\"numRecords\":5000000,\"valueSizeBytes\":1000,\"minutes\":60,\"pipelineTimeout\":240,\"runner\":\"DataflowRunner\"}",
+                  "{\"rowsPerSecond\":50000,\"numRecords\":5000000,\"valueSizeBytes\":1000,\"minutes\":60,\"pipelineTimeout\":180,\"runner\":\"DataflowRunner\"}",
                   Configuration.class));
     } catch (IOException e) {
       throw new RuntimeException(e);
@@ -178,6 +178,13 @@ public void testWriteAndRead() throws IOException, ParseException, InterruptedEx
     PipelineLauncher.LaunchInfo readInfo = readData();
 
     try {
+      // Add monitoring for write job progress
+      PipelineOperator.Result writeResult =
+          pipelineOperator.waitUntilDone(
+              createConfig(writeInfo, Duration.ofMinutes(configuration.pipelineTimeout)));
+      assertNotEquals(PipelineOperator.Result.LAUNCH_FAILED, writeResult);
+      
+      // Add monitoring for read job progress
       PipelineOperator.Result readResult =
           pipelineOperator.waitUntilDone(
               createConfig(readInfo, Duration.ofMinutes(configuration.pipelineTimeout)));
@@ -271,8 +278,12 @@ private PipelineLauncher.LaunchInfo generateDataAndWrite() throws IOException {
             .withProducerConfigUpdates(
                 ImmutableMap.of(
                     ProducerConfig.RETRIES_CONFIG, 10,
-                    ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, 600000,
-                    ProducerConfig.RETRY_BACKOFF_MS_CONFIG, 5000))
+                    ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, 300000, // Reduced from 600000
+                    ProducerConfig.RETRY_BACKOFF_MS_CONFIG, 5000,
+                    ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, 300000, // Add delivery timeout
+                    ProducerConfig.BATCH_SIZE_CONFIG, 16384, // Add batch size
+                    ProducerConfig.LINGER_MS_CONFIG, 100, // Add linger time
+                    ProducerConfig.BUFFER_MEMORY_CONFIG, 33554432)) // Add buffer memory
             .values());
 
     PipelineLauncher.LaunchConfig options =
@@ -287,6 +298,10 @@ private PipelineLauncher.LaunchInfo generateDataAndWrite() throws IOException {
             .addParameter("numWorkers", String.valueOf(configuration.numWorkers))
             .addParameter("maxNumWorkers", String.valueOf(configuration.maxNumWorkers))
             .addParameter("experiments", configuration.useDataflowRunnerV2 ? "use_runner_v2" : "")
+            .addParameter("enableStreamingEngine", "true") // Enable streaming engine
+            .addParameter("streamingMode", "true") // Enable streaming mode
+            .addParameter("usePublicIps", "false") // Use private IPs for better performance
+            .addParameter("subnetwork", "regions/us-central1/subnetworks/default") // Use default subnet
             .build();
 
     return pipelineLauncher.launch(project, region, options);
@@ -298,7 +313,13 @@ private PipelineLauncher.LaunchInfo readData() throws IOException {
         KafkaIO.readBytes()
             .withBootstrapServers(configuration.bootstrapServers)
             .withTopic(kafkaTopic)
-            .withConsumerConfigUpdates(ImmutableMap.of("auto.offset.reset", "earliest"));
+            .withConsumerConfigUpdates(ImmutableMap.of(
+                "auto.offset.reset", "earliest",
+                "session.timeout.ms", "30000", // Add session timeout
+                "heartbeat.interval.ms", "10000", // Add heartbeat interval
+                "max.poll.interval.ms", "300000", // Add max poll interval
+                "fetch.min.bytes", "1", // Add fetch min bytes
+                "fetch.max.wait.ms", "500")); // Add fetch max wait
 
     readPipeline
         .apply("Read from Kafka", readFromKafka)
@@ -311,6 +332,10 @@ private PipelineLauncher.LaunchInfo readData() throws IOException {
             .addParameter("numWorkers", String.valueOf(configuration.numWorkers))
             .addParameter("runner", configuration.runner)
             .addParameter("experiments", configuration.useDataflowRunnerV2 ? "use_runner_v2" : "")
+            .addParameter("enableStreamingEngine", "true") // Enable streaming engine
+            .addParameter("streamingMode", "true") // Enable streaming mode
+            .addParameter("usePublicIps", "false") // Use private IPs for better performance
+            .addParameter("subnetwork", "regions/us-central1/subnetworks/default") // Use default subnet
             .build();
 
     return pipelineLauncher.launch(project, region, options);