checkpoint

aklivity · jfallows · Feb 6, 2024 · Jan 10, 2024 · Jan 11, 2024 · Jan 9, 2024
commit 519527169574e5cc057c2561c331e4e2e3e06afe
diff --git a/...io/aklivity/zilla/runtime/binding/mqtt/kafka/internal/stream/MqttKafkaPublishFactory.java b/...io/aklivity/zilla/runtime/binding/mqtt/kafka/internal/stream/MqttKafkaPublishFactory.java
@@ -64,10 +64,8 @@
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.KafkaResetExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttBeginExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttDataExFW;
-import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttFlushExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttPublishBeginExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttPublishDataExFW;
-import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttPublishFlushExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttResetExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.ResetFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.WindowFW;
@@ -122,7 +120,6 @@ public class MqttKafkaPublishFactory implements MqttKafkaStreamFactory
 
     private final ExtensionFW extensionRO = new ExtensionFW();
     private final MqttBeginExFW mqttBeginExRO = new MqttBeginExFW();
-    private final MqttFlushExFW mqttFlushExRO = new MqttFlushExFW();
     private final MqttDataExFW mqttDataExRO = new MqttDataExFW();
     private final KafkaResetExFW kafkaResetExRO = new KafkaResetExFW();
     private final KafkaBeginExFW kafkaBeginExRO = new KafkaBeginExFW();
@@ -132,7 +129,6 @@ public class MqttKafkaPublishFactory implements MqttKafkaStreamFactory
     private final KafkaFlushExFW.Builder kafkaFlushExRW = new KafkaFlushExFW.Builder();
     private final KafkaDataExFW.Builder kafkaDataExRW = new KafkaDataExFW.Builder();
     private final MqttResetExFW.Builder mqttResetExRW = new MqttResetExFW.Builder();
-    private final MqttFlushExFW.Builder mqttFlushExRW = new MqttFlushExFW.Builder();
     private final Array32FW.Builder<KafkaHeaderFW.Builder, KafkaHeaderFW> kafkaHeadersRW =
         new Array32FW.Builder<>(new KafkaHeaderFW.Builder(), new KafkaHeaderFW());
 
@@ -338,10 +334,6 @@ private void onMqttMessage(
                 final BeginFW begin = beginRO.wrap(buffer, index, index + length);
                 onMqttBegin(begin);
                 break;
-            case FlushFW.TYPE_ID:
-                final FlushFW flush = flushRO.wrap(buffer, index, index + length);
-                onMqttFlush(flush);
-                break;
             case DataFW.TYPE_ID:
                 final DataFW data = dataRO.wrap(buffer, index, index + length);
                 onMqttData(data);
@@ -483,43 +475,6 @@ private String clientHashKey(
             return clientHashKey;
         }
 
-        private void onMqttFlush(
-            FlushFW flush)
-        {
-            final long sequence = flush.sequence();
-            final long acknowledge = flush.acknowledge();
-            final long traceId = flush.traceId();
-            final long authorization = flush.authorization();
-
-            assert acknowledge <= sequence;
-            assert sequence >= initialSeq;
-            assert acknowledge >= initialAck;
-
-            initialSeq = sequence;
-
-            assert initialAck <= initialSeq;
-
-            final OctetsFW extension = flush.extension();
-            final MqttFlushExFW mqttFlushEx = extension.get(mqttFlushExRO::tryWrap);
-
-            assert mqttFlushEx.kind() == MqttFlushExFW.KIND_PUBLISH;
-            final MqttPublishFlushExFW mqttPublishFlushEx = mqttFlushEx.publish();
-
-            final int packetId = mqttPublishFlushEx.packetId();
-
-            final MqttKafkaBindingConfig binding = supplyBinding.apply(routedId);
-            MqttKafkaSessionFactory.MqttSessionProxy session = binding.sessions.get(affinity);
-
-            session.commitOffsetComplete(traceId, authorization, messages.topicString,
-                messages.qos2PartitionId, packetId, this);
-
-            if (hasPublishFlagRetained(publishFlags))
-            {
-                session.commitOffsetComplete(traceId, authorization, retained.topicString,
-                    retained.qos2PartitionId, packetId, this);
-            }
-        }
-
         private void onMqttData(
             DataFW data)
         {
@@ -653,9 +608,8 @@ private void onMqttData(
 
             if ((flags & DATA_FLAG_FIN) != 0x00 && qos == MqttQoS.EXACTLY_ONCE.value())
             {
-                //TODO: when the ack comes for this we need to send the FIN bit, not the
                 session.commitOffsetIncomplete(traceId, authorization, messages.topicString,
-                    messages.qos2PartitionId, packetId, messages);
+                    messages.qos2PartitionId, packetId, messages, false);
             }
 
             if (retainAvailable)
@@ -695,7 +649,7 @@ private void onMqttData(
                     if ((flags & DATA_FLAG_FIN) != 0x00 && qos == MqttQoS.EXACTLY_ONCE.value())
                     {
                         session.commitOffsetIncomplete(traceId, authorization, retained.topicString,
-                            retained.qos2PartitionId, packetId, retained);
+                            retained.qos2PartitionId, packetId, retained, true);
                     }
                 }
                 else
@@ -864,30 +818,6 @@ private void doMqttBegin(
                 traceId, authorization, affinity);
         }
 
-        public void doMqttFlush(
-            long traceId,
-            long authorization,
-            long budgetId,
-            int reserved,
-            int packetId)
-        {
-            if (!hasPublishFlagRetained(publishFlags))
-            {
-                final MqttFlushExFW mqttFlushEx =
-                    mqttFlushExRW.wrap(extBuffer, FlushFW.FIELD_OFFSET_EXTENSION, extBuffer.capacity())
-                        .typeId(mqttTypeId)
-                        .publish(p -> p.packetId(packetId))
-                        .build();
-
-                doFlush(mqtt, originId, routedId, replyId, replySeq, replyAck, replyMax, traceId, authorization,
-                    budgetId, reserved, mqttFlushEx);
-            }
-            else
-            {
-                publishFlags = 0;
-            }
-        }
-
         private void doMqttAbort(
             long traceId,
             long authorization)

diff --git a/...io/aklivity/zilla/runtime/binding/mqtt/kafka/internal/stream/MqttKafkaSessionFactory.java b/...io/aklivity/zilla/runtime/binding/mqtt/kafka/internal/stream/MqttKafkaSessionFactory.java
@@ -99,12 +99,14 @@
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.KafkaTopicPartitionOffsetFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttBeginExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttDataExFW;
+import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttFlushExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttPublishOffsetMetadataFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttResetExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttServerCapabilities;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttSessionBeginExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttSessionDataExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttSessionDataKind;
+import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.MqttSessionFlushExFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.ResetFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.SignalFW;
 import io.aklivity.zilla.runtime.binding.mqtt.kafka.internal.types.stream.WindowFW;
@@ -208,6 +210,7 @@ public class MqttKafkaSessionFactory implements MqttKafkaStreamFactory
 
     private final ExtensionFW extensionRO = new ExtensionFW();
     private final MqttBeginExFW mqttBeginExRO = new MqttBeginExFW();
+    private final MqttFlushExFW mqttFlushExRO = new MqttFlushExFW();
     private final MqttSessionStateFW mqttSessionStateRO = new MqttSessionStateFW();
     private final MqttSessionSignalFW mqttSessionSignalRO = new MqttSessionSignalFW();
     private final MqttWillMessageFW mqttWillRO = new MqttWillMessageFW();
@@ -216,6 +219,7 @@ public class MqttKafkaSessionFactory implements MqttKafkaStreamFactory
     private final MqttPublishOffsetMetadataFW mqttOffsetMetadataRO = new MqttPublishOffsetMetadataFW();
 
     private final MqttResetExFW.Builder mqttResetExRW = new MqttResetExFW.Builder();
+    private final MqttFlushExFW.Builder mqttFlushExRW = new MqttFlushExFW.Builder();
     private final KafkaBeginExFW kafkaBeginExRO = new KafkaBeginExFW();
     private final KafkaDataExFW kafkaDataExRO = new KafkaDataExFW();
     private final KafkaResetExFW kafkaResetExRO = new KafkaResetExFW();
@@ -402,6 +406,8 @@ public final class MqttSessionProxy
         private final String16FW sessionsTopic;
         private final Object2ObjectHashMap<String16FW, MqttKafkaPublishFactory.KafkaMessagesProxy> qos2Publishes;
         private final Long2ObjectHashMap<PublishOffsetMetadata> offsets;
+        private final Int2ObjectHashMap<KafkaTopicPartition> partitions;
+        private final Int2ObjectHashMap<KafkaTopicPartition> retainedPartitions;
         private final List<KafkaPartition> initialisablePartitions;
         private final Long2LongHashMap leaderEpochs;
 
@@ -461,6 +467,8 @@ private MqttSessionProxy(
             this.qos2Publishes = new Object2ObjectHashMap<>();
             this.offsetFetches = new ArrayList<>();
             this.offsets = new Long2ObjectHashMap<>();
+            this.partitions = new Int2ObjectHashMap<>();
+            this.retainedPartitions = new Int2ObjectHashMap<>();
             this.initialisablePartitions = new ArrayList<>();
         }
 
@@ -476,6 +484,10 @@ private void onMqttMessage(
                 final BeginFW begin = beginRO.wrap(buffer, index, index + length);
                 onMqttBegin(begin);
                 break;
+            case FlushFW.TYPE_ID:
+                final FlushFW flush = flushRO.wrap(buffer, index, index + length);
+                onMqttFlush(flush);
+                break;
             case DataFW.TYPE_ID:
                 final DataFW data = dataRO.wrap(buffer, index, index + length);
                 onMqttData(data);
@@ -552,6 +564,40 @@ private void onMqttBegin(
             }
         }
 
+        private void onMqttFlush(
+            FlushFW flush)
+        {
+            final long sequence = flush.sequence();
+            final long acknowledge = flush.acknowledge();
+            final long traceId = flush.traceId();
+            final long authorization = flush.authorization();
+
+            assert acknowledge <= sequence;
+            assert sequence >= initialSeq;
+            assert acknowledge >= initialAck;
+
+            initialSeq = sequence;
+
+            assert initialAck <= initialSeq;
+
+            final OctetsFW extension = flush.extension();
+            final MqttFlushExFW mqttFlushEx = extension.get(mqttFlushExRO::tryWrap);
+
+            assert mqttFlushEx.kind() == MqttFlushExFW.KIND_SESSION;
+            final MqttSessionFlushExFW mqttPublishFlushEx = mqttFlushEx.session();
+
+            final int packetId = mqttPublishFlushEx.packetId();
+
+            final KafkaTopicPartition partition = partitions.remove(packetId);
+            commitOffsetComplete(traceId, authorization, partition.topic, partition.partitionId, packetId);
+
+            if (retainedPartitions.containsKey(packetId))
+            {
+                final KafkaTopicPartition retainedPartition = retainedPartitions.remove(packetId);
+                commitOffsetComplete(traceId, authorization, retainedPartition.topic, retainedPartition.partitionId, packetId);
+            }
+        }
+
         private void onMqttData(
             DataFW data)
         {
@@ -762,7 +808,8 @@ public void commitOffsetIncomplete(
             String topic,
             int partitionId,
             int packetId,
-            MqttKafkaPublishFactory.KafkaProxy kafka)
+            MqttKafkaPublishFactory.KafkaProxy kafka,
+            boolean retained)
         {
             final long offsetKey = offsetKey(topic, partitionId);
             final PublishOffsetMetadata metadata = offsets.get(offsetKey);
@@ -782,6 +829,14 @@ public void commitOffsetIncomplete(
 
             offsetCommit.unfinishedKafkas.add(kafka);
             offsetCommit.unackedPacketIds.add(INCOMPLETE_PACKET_ID);
+            if (retained)
+            {
+                retainedPartitions.put(packetId, new KafkaTopicPartition(topic, partitionId));
+            }
+            else
+            {
+                partitions.put(packetId, new KafkaTopicPartition(topic, partitionId));
+            }
             offsetCommit.doKafkaData(traceId, authorization, 0, DATA_FLAG_COMPLETE, offsetCommitEx);
         }
 
@@ -790,8 +845,7 @@ public void commitOffsetComplete(
             long authorization,
             String topic,
             int partitionId,
-            int packetId,
-            MqttKafkaPublishFactory.MqttPublishProxy publish)
+            int packetId)
         {
             final long offsetKey = offsetKey(topic, partitionId);
             final PublishOffsetMetadata metadata = offsets.get(offsetKey);
@@ -809,7 +863,6 @@ public void commitOffsetComplete(
                     .leaderEpoch((int) leaderEpochs.get(offsetKey)))
                 .build();
 
-            offsetCommit.unackedPublishes.add(publish);
             offsetCommit.unackedPacketIds.add(packetId);
             offsetCommit.doKafkaData(traceId, authorization, 0, DATA_FLAG_COMPLETE, offsetCommitEx);
             metadata.sequence++;
@@ -1127,6 +1180,26 @@ private void doMqttData(
             assert replySeq <= replyAck + replyMax;
         }
 
+        public void doMqttFlush(
+            long traceId,
+            long authorization,
+            long budgetId,
+            int reserved,
+            int packetId)
+        {
+            if (!partitions.containsKey(packetId) && !retainedPartitions.containsKey(packetId))
+            {
+                final MqttFlushExFW mqttFlushEx =
+                    mqttFlushExRW.wrap(extBuffer, FlushFW.FIELD_OFFSET_EXTENSION, extBuffer.capacity())
+                        .typeId(mqttTypeId)
+                        .session(p -> p.packetId(packetId))
+                        .build();
+
+                doFlush(mqtt, originId, routedId, replyId, replySeq, replyAck, replyMax, traceId, authorization,
+                    budgetId, reserved, mqttFlushEx);
+            }
+        }
+
         private void doMqttAbort(
             long traceId,
             long authorization)
@@ -3960,7 +4033,9 @@ private void onKafkaData(
                     {
                         final long offset = partition.partitionOffset();
                         final String16FW metadata = partition.metadata();
-                        final long offsetKey = offsetKey(topic, partition.partitionId());
+                        final int partitionId = partition.partitionId();
+                        final long offsetKey = offsetKey(topic, partitionId);
+
                         delegate.leaderEpochs.put(offsetKey, partition.leaderEpoch());
 
                         PublishOffsetMetadata offsetMetadata;
@@ -3977,6 +4052,17 @@ private void onKafkaData(
                                 delegate.offsetCommit.doKafkaBegin(traceId, authorization, 0);
                             }
                             delegate.offsets.put(offsetKey, offsetMetadata);
+                            if (!retained)
+                            {
+                                offsetMetadata.packetIds.forEach(p ->
+                                    delegate.partitions.put(p, new KafkaTopicPartition(topic, partitionId)));
+                            }
+                            else
+                            {
+                                offsetMetadata.packetIds.forEach(p ->
+                                    delegate.retainedPartitions.put(p, new KafkaTopicPartition(topic, partitionId)));
+                            }
+
                         }
                         else
                         {
@@ -4515,8 +4601,19 @@ else if (wasOpen)
                     }
                     else
                     {
-                        final MqttKafkaPublishFactory.MqttPublishProxy publish = unackedPublishes.remove();
-                        publish.doMqttFlush(traceId, authorization, 0, 0, packetId);
+                        if (delegate.partitions.containsKey(packetId))
+                        {
+                            delegate.partitions.remove(packetId);
+                        }
+                        else if (delegate.retainedPartitions.containsKey(packetId))
+                        {
+                            delegate.retainedPartitions.remove(packetId);
+                        }
+
+                        if (!delegate.partitions.containsKey(packetId) && !delegate.retainedPartitions.containsKey(packetId))
+                        {
+                            delegate.doMqttFlush(traceId, authorization, 0, 0, packetId);
+                        }
                     }
                 }
                 else
@@ -4676,6 +4773,20 @@ private KafkaPartition(
         }
     }
 
+    public static final class KafkaTopicPartition
+    {
+        private final String topic;
+        private final int partitionId;
+
+        KafkaTopicPartition(
+            String topic,
+            int partitionId)
+        {
+            this.topic = topic;
+            this.partitionId = partitionId;
+        }
+    }
+
     public static final class PublishOffsetMetadata
     {
         public final long producerId;