KAFKA-21307: Add share group DLQ setup to ShareConsumeBenchSpec

trogdor/src/main/java/org/apache/kafka/trogdor/workload/RandomRecordAcknowledger.java

@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.trogdor.workload;
+
+import org.apache.kafka.clients.consumer.AcknowledgeType;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.trogdor.common.JsonUtil;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.JsonNode;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * This class acknowledges each record with a randomly chosen type. The percentage of records that are rejected and
+ * released is configurable, and all the remaining records are accepted. The number of records acknowledged with each
+ * type is present in the status from the `ShareConsumeBenchWorker` class.
+ *
+ * Example spec:
+ * {
+ *    "type": "random",
+ *    "rejectPercent": 20,
+ *    "releasePercent": 10
+ * }
+ *
+ * This will reject 20% of the records, release 10% of them and accept the remaining 70%.
+ */
+public class RandomRecordAcknowledger implements RecordAcknowledger {
+    private final int rejectPercent;
+    private final int releasePercent;
+    private final AtomicLong accepted = new AtomicLong();
+    private final AtomicLong rejected = new AtomicLong();
+    private final AtomicLong released = new AtomicLong();
+
+    @JsonCreator
+    public RandomRecordAcknowledger(@JsonProperty("rejectPercent") int rejectPercent,
+                                    @JsonProperty("releasePercent") int releasePercent) {
+        if (rejectPercent < 0 || releasePercent < 0 || rejectPercent + releasePercent > 100) {
+            throw new IllegalArgumentException(String.format(
+                    "rejectPercent and releasePercent must be non-negative and sum to at most 100, but were %d and %d",
+                    rejectPercent, releasePercent));
+        }
+        this.rejectPercent = rejectPercent;
+        this.releasePercent = releasePercent;
+    }
+
+    @JsonProperty
+    public int rejectPercent() {
+        return rejectPercent;
+    }
+
+    @JsonProperty
+    public int releasePercent() {
+        return releasePercent;
+    }
+
+    @Override
+    public AcknowledgeType acknowledge(ConsumerRecord<byte[], byte[]> record) {
+        int roll = ThreadLocalRandom.current().nextInt(100);
+        if (roll < rejectPercent) {
+            rejected.incrementAndGet();
+            return AcknowledgeType.REJECT;
+        } else if (roll < rejectPercent + releasePercent) {
+            released.incrementAndGet();
+            return AcknowledgeType.RELEASE;
+        }
+        accepted.incrementAndGet();
+        return AcknowledgeType.ACCEPT;
+    }
+
+    @Override
+    public JsonNode acknowledgerStatus() {
+        Map<String, Long> counts = new LinkedHashMap<>();
+        counts.put("accepted", accepted.get());
+        counts.put("rejected", rejected.get());
+        counts.put("released", released.get());
+        return JsonUtil.JSON_SERDE.valueToTree(counts);
+    }
+}

trogdor/src/main/java/org/apache/kafka/trogdor/workload/RecordAcknowledger.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.trogdor.workload;
+
+import org.apache.kafka.clients.consumer.AcknowledgeType;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import com.fasterxml.jackson.databind.JsonNode;
+
+/**
+ * RecordAcknowledger decides how each record consumed by a ShareConsumeBench workload is acknowledged.
+ * A single instance is shared by all the share consumers of the worker, so implementations must be thread safe.
+ */
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
+@JsonSubTypes(value = {
+    @JsonSubTypes.Type(value = RandomRecordAcknowledger.class, name = "random"),
+})
+public interface RecordAcknowledger {
+    /**
+     * Returns the acknowledge type to apply to the given record. It must not be null.
+     */
+    AcknowledgeType acknowledge(ConsumerRecord<byte[], byte[]> record);
+    JsonNode acknowledgerStatus();
+}

trogdor/src/main/java/org/apache/kafka/trogdor/workload/ShareConsumeBenchSpec.java

@@ -66,6 +66,14 @@
  * share consumers use implicit acknowledgement. Note that with "RELEASE" the records are redelivered, so the
  * "maxMessages" field bounds the total number of deliveries rather than the number of distinct records.
  *
+ * The optional "recordAcknowledger" field is like "acknowledgeType", but lets the acknowledge type be decided per record
+ * by a `RecordAcknowledger`. It also makes the share consumers use explicit acknowledgement mode, and cannot be
+ * combined with "acknowledgeType".
+ *
+ * The optional "shareGroupDLQConfig" field sets up a dead-letter queue for the share group before the consumers start.
+ * See `ShareGroupDLQConfig`. Records rejected, or released until their delivery attempts are exhausted, are then written
+ * to the DLQ topic.
+ *
  * An example JSON representation which will result in a share consumer that is part of the share group "sg" and
  * subscribed to topics foo1, foo2, foo3 and bar.
  * #{@code
@@ -96,6 +104,8 @@ public final class ShareConsumeBenchSpec extends TaskSpec {
     private final int threadsPerWorker;
     private final Optional<RecordProcessor> recordProcessor;
     private final Optional<AcknowledgeType> acknowledgeType;
+    private final Optional<RecordAcknowledger> recordAcknowledger;
+    private final Optional<ShareGroupDLQConfig> shareGroupDLQConfig;
 
     @JsonCreator
     public ShareConsumeBenchSpec(@JsonProperty("startMs") long startMs,
@@ -111,7 +121,9 @@ public ShareConsumeBenchSpec(@JsonProperty("startMs") long startMs,
                             @JsonProperty("threadsPerWorker") Integer threadsPerWorker,
                             @JsonProperty("recordProcessor") Optional<RecordProcessor> recordProcessor,
                             @JsonProperty("activeTopics") List<String> activeTopics,
-                            @JsonProperty("acknowledgeType") AcknowledgeType acknowledgeType) {
+                            @JsonProperty("acknowledgeType") AcknowledgeType acknowledgeType,
+                            @JsonProperty("recordAcknowledger") Optional<RecordAcknowledger> recordAcknowledger,
+                            @JsonProperty("shareGroupDLQConfig") ShareGroupDLQConfig shareGroupDLQConfig) {
         super(startMs, durationMs);
         this.consumerNode = (consumerNode == null) ? "" : consumerNode;
         this.bootstrapServers = (bootstrapServers == null) ? "" : bootstrapServers;
@@ -124,10 +136,20 @@ public ShareConsumeBenchSpec(@JsonProperty("startMs") long startMs,
         this.shareGroup = shareGroup == null ? DEFAULT_SHARE_GROUP_NAME : shareGroup;
         this.threadsPerWorker = threadsPerWorker == null ? 1 : threadsPerWorker;
         this.recordProcessor = recordProcessor;
-        if (acknowledgeType == AcknowledgeType.RENEW) {
-            throw new IllegalArgumentException("acknowledgeType must be ACCEPT, REJECT or RELEASE, but was " + acknowledgeType);
-        }
         this.acknowledgeType = Optional.ofNullable(acknowledgeType);
+        this.recordAcknowledger = recordAcknowledger == null ? Optional.empty() : recordAcknowledger;
+        this.shareGroupDLQConfig = Optional.ofNullable(shareGroupDLQConfig);
+        validateAcknowledgement(this.acknowledgeType, this.recordAcknowledger);
+    }
+
+    private static void validateAcknowledgement(Optional<AcknowledgeType> acknowledgeType,
+                                                Optional<RecordAcknowledger> recordAcknowledger) {
+        if (acknowledgeType.filter(type -> type == AcknowledgeType.RENEW).isPresent()) {
+            throw new IllegalArgumentException("acknowledgeType must be ACCEPT, REJECT or RELEASE, but was " + AcknowledgeType.RENEW);
+        }
+        if (acknowledgeType.isPresent() && recordAcknowledger.isPresent()) {
+            throw new IllegalArgumentException("acknowledgeType and recordAcknowledger cannot both be set");
+        }
     }
 
     @JsonProperty
@@ -170,6 +192,16 @@ public Optional<AcknowledgeType> acknowledgeType() {
         return acknowledgeType;
     }
 
+    @JsonProperty
+    public Optional<RecordAcknowledger> recordAcknowledger() {
+        return recordAcknowledger;
+    }
+
+    @JsonProperty
+    public Optional<ShareGroupDLQConfig> shareGroupDLQConfig() {
+        return shareGroupDLQConfig;
+    }
+
     @JsonProperty
     public Map<String, String> consumerConf() {
         return consumerConf;