This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new fea5bd2490a Fix oversized subscription requests with slicing (#18695)
fea5bd2490a is described below
commit fea5bd2490a8274cebf6fb197744d49321859de2
Author: Caideyipi <[email protected]>
AuthorDate: Mon Oct 12 10:39:19 2026 +0800
Fix oversized subscription requests with slicing (#18695)
---
.../subscription/i18n/SubscriptionMessages.java | 4 +
.../subscription/i18n/SubscriptionMessages.java | 4 +
.../payload/request/PipeSubscribeRequestType.java | 1 +
.../payload/request/PipeSubscribeSliceReq.java | 169 +++++++++++++++
.../request/PipeSubscribeSliceReqBuilder.java | 80 +++++++
.../request/PipeSubscribeSliceReqHandler.java | 141 +++++++++++++
.../SubscriptionSessionConnection.java | 8 +
.../base/AbstractSubscriptionProvider.java | 186 ++++++++++++++++-
.../payload/request/PipeSubscribeSliceReqTest.java | 203 ++++++++++++++++++
.../base/SubscriptionProviderSlicingTest.java | 232 +++++++++++++++++++++
.../receiver/SubscriptionReceiverV1.java | 52 +++++
.../receiver/SubscriptionReceiverV1Test.java | 55 +++++
12 files changed, 1134 insertions(+), 1 deletion(-)
diff --git
a/iotdb-client/subscription/src/main/i18n/en/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
b/iotdb-client/subscription/src/main/i18n/en/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
index 41bcc525d18..a161eea5151 100644
---
a/iotdb-client/subscription/src/main/i18n/en/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
+++
b/iotdb-client/subscription/src/main/i18n/en/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
@@ -21,6 +21,10 @@ package org.apache.iotdb.rpc.subscription.i18n;
public final class SubscriptionMessages {
+ public static final String
+
EXCEPTION_SUBSCRIPTION_REQUEST_EXCEEDS_THE_NEGOTIATED_FRAME_LIMIT_CA8DEB50 =
+ "Subscription request exceeds the negotiated frame limit";
+
// --- TopicConstant / ConsumerConstant ---
public static final String UTILITY_CLASS = "Utility class";
diff --git
a/iotdb-client/subscription/src/main/i18n/zh/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
b/iotdb-client/subscription/src/main/i18n/zh/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
index 89e4d59750a..75f3c62bdf9 100644
---
a/iotdb-client/subscription/src/main/i18n/zh/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
+++
b/iotdb-client/subscription/src/main/i18n/zh/org/apache/iotdb/rpc/subscription/i18n/SubscriptionMessages.java
@@ -21,6 +21,10 @@ package org.apache.iotdb.rpc.subscription.i18n;
public final class SubscriptionMessages {
+ public static final String
+
EXCEPTION_SUBSCRIPTION_REQUEST_EXCEEDS_THE_NEGOTIATED_FRAME_LIMIT_CA8DEB50 =
+ "订阅请求超过协商的帧大小限制";
+
// --- TopicConstant / ConsumerConstant ---
public static final String UTILITY_CLASS = "工具类";
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeRequestType.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeRequestType.java
index 9fcc1d86b0c..5a6c414695c 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeRequestType.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeRequestType.java
@@ -32,6 +32,7 @@ public enum PipeSubscribeRequestType {
SUBSCRIBE((short) 5),
UNSUBSCRIBE((short) 6),
SEEK((short) 7),
+ SLICE((short) 8),
;
private final short type;
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReq.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReq.java
new file mode 100644
index 00000000000..d783e93bf34
--- /dev/null
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReq.java
@@ -0,0 +1,169 @@
+/*
+ * 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.iotdb.rpc.subscription.payload.request;
+
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
+
+import org.apache.tsfile.utils.Binary;
+import org.apache.tsfile.utils.PublicBAOS;
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Objects;
+
+public class PipeSubscribeSliceReq extends TPipeSubscribeReq {
+
+ private transient int orderId;
+ private transient byte originReqVersion;
+ private transient short originReqType;
+ private transient int originBodySize;
+ private transient int sliceIndex;
+ private transient int sliceCount;
+ private transient byte[] sliceBody;
+
+ public int getOrderId() {
+ return orderId;
+ }
+
+ public byte getOriginReqVersion() {
+ return originReqVersion;
+ }
+
+ public short getOriginReqType() {
+ return originReqType;
+ }
+
+ public int getOriginBodySize() {
+ return originBodySize;
+ }
+
+ public int getSliceIndex() {
+ return sliceIndex;
+ }
+
+ public int getSliceCount() {
+ return sliceCount;
+ }
+
+ public byte[] getSliceBody() {
+ return sliceBody;
+ }
+
+ public static PipeSubscribeSliceReq toTPipeSubscribeReq(
+ final TPipeSubscribeReq originReq,
+ final int orderId,
+ final int sliceIndex,
+ final int sliceCount,
+ final int bodySizeLimit)
+ throws IOException {
+ final ByteBuffer originBody = originReq.body.duplicate();
+ final int originPosition = originBody.position();
+ final int originBodySize = originBody.remaining();
+ final int startOffset = sliceIndex * bodySizeLimit;
+ final int endOffset = startOffset + Math.min(bodySizeLimit, originBodySize
- startOffset);
+
+ final PipeSubscribeSliceReq sliceReq = new PipeSubscribeSliceReq();
+ sliceReq.orderId = orderId;
+ sliceReq.originReqVersion = originReq.getVersion();
+ sliceReq.originReqType = originReq.getType();
+ sliceReq.originBodySize = originBodySize;
+ sliceReq.sliceIndex = sliceIndex;
+ sliceReq.sliceCount = sliceCount;
+ sliceReq.sliceBody = new byte[endOffset - startOffset];
+ originBody.position(originPosition + startOffset);
+ originBody.get(sliceReq.sliceBody);
+
+ sliceReq.version = PipeSubscribeRequestVersion.VERSION_1.getVersion();
+ sliceReq.type = PipeSubscribeRequestType.SLICE.getType();
+ try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+ final DataOutputStream outputStream = new
DataOutputStream(byteArrayOutputStream)) {
+ ReadWriteIOUtils.write(sliceReq.orderId, outputStream);
+ ReadWriteIOUtils.write(sliceReq.originReqVersion, outputStream);
+ ReadWriteIOUtils.write(sliceReq.originReqType, outputStream);
+ ReadWriteIOUtils.write(sliceReq.originBodySize, outputStream);
+ ReadWriteIOUtils.write(sliceReq.sliceIndex, outputStream);
+ ReadWriteIOUtils.write(sliceReq.sliceCount, outputStream);
+ ReadWriteIOUtils.write(new Binary(sliceReq.sliceBody), outputStream);
+ sliceReq.body =
+ ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0,
byteArrayOutputStream.size());
+ }
+ return sliceReq;
+ }
+
+ public static PipeSubscribeSliceReq fromTPipeSubscribeReq(final
TPipeSubscribeReq thriftReq) {
+ final PipeSubscribeSliceReq sliceReq = new PipeSubscribeSliceReq();
+ final ByteBuffer body = thriftReq.body.duplicate();
+ sliceReq.orderId = ReadWriteIOUtils.readInt(body);
+ sliceReq.originReqVersion = ReadWriteIOUtils.readByte(body);
+ sliceReq.originReqType = ReadWriteIOUtils.readShort(body);
+ sliceReq.originBodySize = ReadWriteIOUtils.readInt(body);
+ sliceReq.sliceIndex = ReadWriteIOUtils.readInt(body);
+ sliceReq.sliceCount = ReadWriteIOUtils.readInt(body);
+ final int sliceBodySize = ReadWriteIOUtils.readInt(body);
+ if (sliceBodySize < 0 || sliceBodySize != body.remaining()) {
+ throw new IllegalArgumentException();
+ }
+ sliceReq.sliceBody = new byte[sliceBodySize];
+ body.get(sliceReq.sliceBody);
+ sliceReq.version = thriftReq.version;
+ sliceReq.type = thriftReq.type;
+ sliceReq.body = thriftReq.body;
+ return sliceReq;
+ }
+
+ @Override
+ public boolean equals(final Object obj) {
+ if (this == obj) {
+ return true;
+ }
+ if (obj == null || getClass() != obj.getClass()) {
+ return false;
+ }
+ final PipeSubscribeSliceReq that = (PipeSubscribeSliceReq) obj;
+ return orderId == that.orderId
+ && originReqVersion == that.originReqVersion
+ && originReqType == that.originReqType
+ && originBodySize == that.originBodySize
+ && sliceIndex == that.sliceIndex
+ && sliceCount == that.sliceCount
+ && Arrays.equals(sliceBody, that.sliceBody)
+ && version == that.version
+ && type == that.type
+ && Objects.equals(body, that.body);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(
+ orderId,
+ originReqVersion,
+ originReqType,
+ originBodySize,
+ sliceIndex,
+ sliceCount,
+ Arrays.hashCode(sliceBody),
+ version,
+ type,
+ body);
+ }
+}
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqBuilder.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqBuilder.java
new file mode 100644
index 00000000000..2a743cfbebd
--- /dev/null
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqBuilder.java
@@ -0,0 +1,80 @@
+/*
+ * 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.iotdb.rpc.subscription.payload.request;
+
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
+
+import java.io.IOException;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public final class PipeSubscribeSliceReqBuilder {
+
+ private static final AtomicInteger ORDER_ID_GENERATOR = new AtomicInteger(0);
+
+ private PipeSubscribeSliceReqBuilder() {}
+
+ public static boolean shouldSlice(final TPipeSubscribeReq req, final int
bodySizeLimit) {
+ return bodySizeLimit > 0
+ && Objects.nonNull(req)
+ && req.getVersion() ==
PipeSubscribeRequestVersion.VERSION_1.getVersion()
+ && req.isSetBody()
+ && req.body.remaining() > bodySizeLimit
+ && isSliceableRequestType(req.getType());
+ }
+
+ public static boolean isSliceableRequestType(final short type) {
+ return type != PipeSubscribeRequestType.HANDSHAKE.getType()
+ && type != PipeSubscribeRequestType.CLOSE.getType()
+ && type != PipeSubscribeRequestType.SLICE.getType()
+ && PipeSubscribeRequestType.isValidatedRequestType(type);
+ }
+
+ public static int nextOrderId() {
+ return ORDER_ID_GENERATOR.getAndIncrement();
+ }
+
+ public static int getSliceCount(final TPipeSubscribeReq req, final int
bodySizeLimit) {
+ if (bodySizeLimit <= 0 || Objects.isNull(req) || !req.isSetBody()) {
+ throw new IllegalArgumentException();
+ }
+ final int bodySize = req.body.remaining();
+ return bodySize == 0 ? 1 : (bodySize - 1) / bodySizeLimit + 1;
+ }
+
+ public static PipeSubscribeSliceReq buildSliceReq(
+ final TPipeSubscribeReq originReq,
+ final int orderId,
+ final int sliceIndex,
+ final int sliceCount,
+ final int bodySizeLimit)
+ throws IOException {
+ if (bodySizeLimit <= 0
+ || sliceCount <= 1
+ || sliceIndex < 0
+ || sliceIndex >= sliceCount
+ || !shouldSlice(originReq, bodySizeLimit)
+ || getSliceCount(originReq, bodySizeLimit) != sliceCount) {
+ throw new IllegalArgumentException();
+ }
+ return PipeSubscribeSliceReq.toTPipeSubscribeReq(
+ originReq, orderId, sliceIndex, sliceCount, bodySizeLimit);
+ }
+}
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqHandler.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqHandler.java
new file mode 100644
index 00000000000..e62983aa4c6
--- /dev/null
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqHandler.java
@@ -0,0 +1,141 @@
+/*
+ * 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.iotdb.rpc.subscription.payload.request;
+
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
+
+import org.apache.tsfile.utils.PublicBAOS;
+
+import java.nio.ByteBuffer;
+import java.util.Objects;
+import java.util.Optional;
+
+public class PipeSubscribeSliceReqHandler {
+
+ private static final int INITIAL_BUFFER_SIZE = 8 * 1024;
+
+ private final int maxOriginBodySize;
+
+ private int orderId;
+ private byte originReqVersion = -1;
+ private short originReqType = -1;
+ private int originBodySize = -1;
+ private int sliceCount = -1;
+ private int nextSliceIndex;
+ private int receivedBodySize;
+ private PublicBAOS assembledBody;
+
+ public PipeSubscribeSliceReqHandler(final int maxOriginBodySize) {
+ if (maxOriginBodySize <= 0) {
+ throw new IllegalArgumentException();
+ }
+ this.maxOriginBodySize = maxOriginBodySize;
+ }
+
+ public synchronized boolean receiveSlice(final PipeSubscribeSliceReq req) {
+ if (!isValidSlice(req)) {
+ clear();
+ return false;
+ }
+
+ if (assembledBody == null) {
+ if (req.getSliceIndex() != 0) {
+ return false;
+ }
+ orderId = req.getOrderId();
+ originReqVersion = req.getOriginReqVersion();
+ originReqType = req.getOriginReqType();
+ originBodySize = req.getOriginBodySize();
+ sliceCount = req.getSliceCount();
+ assembledBody = new PublicBAOS(Math.min(INITIAL_BUFFER_SIZE,
originBodySize));
+ }
+
+ if (orderId != req.getOrderId()
+ || originReqVersion != req.getOriginReqVersion()
+ || originReqType != req.getOriginReqType()
+ || originBodySize != req.getOriginBodySize()
+ || sliceCount != req.getSliceCount()
+ || nextSliceIndex != req.getSliceIndex()) {
+ clear();
+ return false;
+ }
+
+ final byte[] sliceBody = req.getSliceBody();
+ if (receivedBodySize > originBodySize - sliceBody.length) {
+ clear();
+ return false;
+ }
+ assembledBody.write(sliceBody, 0, sliceBody.length);
+ receivedBodySize += sliceBody.length;
+ nextSliceIndex++;
+
+ if (nextSliceIndex < sliceCount && receivedBodySize >= originBodySize) {
+ clear();
+ return false;
+ }
+ if (nextSliceIndex == sliceCount && receivedBodySize != originBodySize) {
+ clear();
+ return false;
+ }
+ return true;
+ }
+
+ private boolean isValidSlice(final PipeSubscribeSliceReq req) {
+ return Objects.nonNull(req)
+ && req.getVersion() ==
PipeSubscribeRequestVersion.VERSION_1.getVersion()
+ && req.getType() == PipeSubscribeRequestType.SLICE.getType()
+ && req.getOriginReqVersion() ==
PipeSubscribeRequestVersion.VERSION_1.getVersion()
+ &&
PipeSubscribeSliceReqBuilder.isSliceableRequestType(req.getOriginReqType())
+ && req.getOriginBodySize() > 0
+ && req.getOriginBodySize() <= maxOriginBodySize
+ && req.getSliceCount() > 1
+ && req.getSliceIndex() >= 0
+ && req.getSliceIndex() < req.getSliceCount()
+ && Objects.nonNull(req.getSliceBody())
+ && req.getSliceBody().length > 0
+ && req.getSliceBody().length <= req.getOriginBodySize();
+ }
+
+ public synchronized Optional<TPipeSubscribeReq> makeReqIfComplete() {
+ if (assembledBody == null
+ || nextSliceIndex != sliceCount
+ || receivedBodySize != originBodySize) {
+ return Optional.empty();
+ }
+
+ final TPipeSubscribeReq req = new TPipeSubscribeReq();
+ req.version = originReqVersion;
+ req.type = originReqType;
+ req.body = ByteBuffer.wrap(assembledBody.getBuf(), 0,
assembledBody.size());
+ clear();
+ return Optional.of(req);
+ }
+
+ public synchronized void clear() {
+ orderId = 0;
+ originReqVersion = -1;
+ originReqType = -1;
+ originBodySize = -1;
+ sliceCount = -1;
+ nextSliceIndex = 0;
+ receivedBodySize = 0;
+ assembledBody = null;
+ }
+}
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
index df3de1331e9..baff5f6034a 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
@@ -21,6 +21,7 @@ package org.apache.iotdb.session.subscription;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.service.rpc.thrift.ServerProperties;
import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp;
import org.apache.iotdb.session.Session;
@@ -59,6 +60,13 @@ public class SubscriptionSessionConnection extends
SessionConnection {
return client.pipeSubscribe(req);
}
+ public int getServerThriftMaxFrameSize(final int fallbackMaxFrameSize)
throws TException {
+ final ServerProperties properties = client.getProperties();
+ return properties.isSetThriftMaxFrameSize()
+ ? properties.getThriftMaxFrameSize()
+ : fallbackMaxFrameSize;
+ }
+
public boolean setTimeout(final int timeoutInMs) {
return setTransportTimeout(timeoutInMs);
}
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
index 477a7ae45cb..db0ded818ec 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.session.subscription.consumer.base;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.rpc.subscription.config.ConsumerConfig;
import org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
import org.apache.iotdb.rpc.subscription.config.TopicConfig;
@@ -36,6 +37,7 @@ import
org.apache.iotdb.rpc.subscription.i18n.SubscriptionMessages;
import org.apache.iotdb.rpc.subscription.payload.poll.PollFilePayload;
import org.apache.iotdb.rpc.subscription.payload.poll.PollPayload;
import org.apache.iotdb.rpc.subscription.payload.poll.PollTabletsPayload;
+import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress;
import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollRequest;
import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollRequestType;
@@ -45,6 +47,7 @@ import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeCloseReq;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeCommitReq;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeHandshakeReq;
import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribePollReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeSliceReqBuilder;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeSubscribeReq;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeUnsubscribeReq;
import
org.apache.iotdb.rpc.subscription.payload.request.SubscriptionHeartbeatReq;
@@ -66,8 +69,11 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -79,6 +85,9 @@ public abstract class AbstractSubscriptionProvider {
private static final Logger LOGGER =
LoggerFactory.getLogger(AbstractSubscriptionProvider.class);
private static final int CLOSE_TIMEOUT_IN_MS = 5_000;
+ private static final int REQUEST_FRAME_RESERVED_BYTES = 1024;
+ // A commit response echoes accepted contexts and may also carry committed
progress.
+ private static final int COMMIT_RESPONSE_HEADROOM_FACTOR = 3;
private static final String STATUS_FORMATTER = "Status code is [%s], status
message is [%s].";
private static final String INTERNAL_ERROR_FORMATTER =
@@ -101,6 +110,9 @@ public abstract class AbstractSubscriptionProvider {
private final AtomicBoolean isAvailable = new AtomicBoolean(false);
private final ReentrantLock rpcLock = new ReentrantLock();
+ private final int localThriftMaxFrameSize;
+ private volatile int requestBodySizeLimit;
+
private final TEndPoint endPoint;
private int dataNodeId;
@@ -149,6 +161,8 @@ public abstract class AbstractSubscriptionProvider {
this.username = username;
this.password = password;
this.encryptedPassword = encryptedPassword;
+ this.localThriftMaxFrameSize = thriftMaxFrameSize;
+ this.requestBodySizeLimit = getRequestBodySizeLimit(thriftMaxFrameSize);
this.heartbeatIntervalMs = heartbeatIntervalMs;
this.connectionTimeoutInMs = connectionTimeoutInMs;
}
@@ -161,12 +175,66 @@ public abstract class AbstractSubscriptionProvider {
throws TException, IoTDBConnectionException {
rpcLock.lock();
try {
- return getSessionConnection().pipeSubscribe(req);
+ final SubscriptionSessionConnection connection = getSessionConnection();
+ return pipeSubscribeWithOptionalSlicing(req, requestBodySizeLimit,
connection::pipeSubscribe);
} finally {
rpcLock.unlock();
}
}
+ static TPipeSubscribeResp pipeSubscribeWithOptionalSlicing(
+ final TPipeSubscribeReq req,
+ final int bodySizeLimit,
+ final PipeSubscribeRequestSender requestSender)
+ throws TException {
+ if (!PipeSubscribeSliceReqBuilder.shouldSlice(req, bodySizeLimit)) {
+ if (bodySizeLimit <= 0 || (req.isSetBody() && req.body.remaining() >
bodySizeLimit)) {
+ throw new TException(
+ SubscriptionMessages
+
.EXCEPTION_SUBSCRIPTION_REQUEST_EXCEEDS_THE_NEGOTIATED_FRAME_LIMIT_CA8DEB50);
+ }
+ return requestSender.send(req);
+ }
+
+ final int orderId = PipeSubscribeSliceReqBuilder.nextOrderId();
+ final int sliceCount = PipeSubscribeSliceReqBuilder.getSliceCount(req,
bodySizeLimit);
+ TPipeSubscribeResp resp = null;
+ for (int sliceIndex = 0; sliceIndex < sliceCount; sliceIndex++) {
+ try {
+ resp =
+ requestSender.send(
+ PipeSubscribeSliceReqBuilder.buildSliceReq(
+ req, orderId, sliceIndex, sliceCount, bodySizeLimit));
+ } catch (final IOException e) {
+ throw new TException(e);
+ }
+ if (sliceIndex < sliceCount - 1
+ && (resp == null
+ || resp.getStatus() == null
+ || resp.getStatus().getCode() !=
TSStatusCode.SUCCESS_STATUS.getStatusCode())) {
+ return resp;
+ }
+ }
+ return resp;
+ }
+
+ private static int getRequestBodySizeLimit(final int maxFrameSize) {
+ return Math.max(0, maxFrameSize - REQUEST_FRAME_RESERVED_BYTES);
+ }
+
+ private void updateRequestBodySizeLimit() {
+ try {
+ final int serverMaxFrameSize =
+
getSessionConnection().getServerThriftMaxFrameSize(localThriftMaxFrameSize);
+ if (serverMaxFrameSize > 0) {
+ requestBodySizeLimit =
+ getRequestBodySizeLimit(Math.min(localThriftMaxFrameSize,
serverMaxFrameSize));
+ }
+ } catch (final TException | IoTDBConnectionException ignored) {
+ // Keep the local limit for servers that do not expose their frame size.
+ }
+ }
+
boolean isAvailable() {
return isAvailable.get();
}
@@ -207,6 +275,7 @@ public abstract class AbstractSubscriptionProvider {
}
session.open(); // throw IoTDBConnectionException
+ updateRequestBodySizeLimit();
// TODO: pass the complete consumer parameter configuration to the server
final Map<String, String> consumerAttributes = new HashMap<>();
@@ -371,6 +440,40 @@ public abstract class AbstractSubscriptionProvider {
PipeSubscribeHeartbeatResp heartbeat(
final List<SubscriptionCommitContext> processorBufferedCommitContexts)
throws SubscriptionException {
+ final List<List<SubscriptionCommitContext>> batches;
+ try {
+ batches =
+ partitionCommitContexts(
+ processorBufferedCommitContexts, requestBodySizeLimit,
Integer.BYTES);
+ } catch (final IOException e) {
+ LOGGER.warn(
+ SubscriptionMessages
+
.LOG_IOEXCEPTION_OCCURRED_SUBSCRIPTIONPROVIDER_ARG_SERIALIZE_HEARTBEAT_REQUEST_ARG_634C8333,
+ this,
+ processorBufferedCommitContexts,
+ e);
+ throw new SubscriptionRuntimeNonCriticalException(e.getMessage(), e);
+ }
+
+ PipeSubscribeHeartbeatResp latestResp = null;
+ final Set<String> topicNamesToUnsubscribe = new LinkedHashSet<>();
+ rpcLock.lock();
+ try {
+ for (final List<SubscriptionCommitContext> batch : batches) {
+ latestResp = heartbeatOnce(batch);
+
topicNamesToUnsubscribe.addAll(latestResp.getTopicNamesToUnsubscribe());
+ }
+ } finally {
+ rpcLock.unlock();
+ }
+ latestResp.getTopicNamesToUnsubscribe().clear();
+ latestResp.getTopicNamesToUnsubscribe().addAll(topicNamesToUnsubscribe);
+ return latestResp;
+ }
+
+ private PipeSubscribeHeartbeatResp heartbeatOnce(
+ final List<SubscriptionCommitContext> processorBufferedCommitContexts)
+ throws SubscriptionException {
final SubscriptionHeartbeatReq req;
try {
req =
SubscriptionHeartbeatReq.toThriftReq(processorBufferedCommitContexts);
@@ -634,6 +737,41 @@ public abstract class AbstractSubscriptionProvider {
CommitResult commit(
final List<SubscriptionCommitContext> subscriptionCommitContexts, final
boolean nack)
throws SubscriptionException {
+ final List<List<SubscriptionCommitContext>> batches;
+ try {
+ batches =
+ partitionCommitContexts(
+ subscriptionCommitContexts,
+ Math.max(1, requestBodySizeLimit /
COMMIT_RESPONSE_HEADROOM_FACTOR),
+ Integer.BYTES + Byte.BYTES);
+ } catch (final IOException e) {
+ LOGGER.warn(
+ SubscriptionMessages
+
.LOG_IOEXCEPTION_OCCURRED_SUBSCRIPTIONPROVIDER_ARG_SERIALIZE_COMMIT_REQUEST_ARG_D5335538,
+ this,
+ subscriptionCommitContexts,
+ e);
+ throw new SubscriptionRuntimeNonCriticalException(e.getMessage(), e);
+ }
+
+ final List<SubscriptionCommitContext> acceptedCommitContexts = new
ArrayList<>();
+ final Map<String, TopicProgress> committedProgressByTopic = new
LinkedHashMap<>();
+ rpcLock.lock();
+ try {
+ for (final List<SubscriptionCommitContext> batch : batches) {
+ final CommitResult result = commitOnce(batch, nack);
+ acceptedCommitContexts.addAll(result.getAcceptedCommitContexts());
+ mergeTopicProgress(committedProgressByTopic,
result.getCommittedProgressByTopic());
+ }
+ } finally {
+ rpcLock.unlock();
+ }
+ return new CommitResult(acceptedCommitContexts, committedProgressByTopic);
+ }
+
+ private CommitResult commitOnce(
+ final List<SubscriptionCommitContext> subscriptionCommitContexts, final
boolean nack)
+ throws SubscriptionException {
final PipeSubscribeCommitReq req;
try {
req =
PipeSubscribeCommitReq.toTPipeSubscribeReq(subscriptionCommitContexts, nack);
@@ -666,6 +804,52 @@ public abstract class AbstractSubscriptionProvider {
commitResp.getAcceptedCommitContexts(),
commitResp.getCommittedProgressByTopic());
}
+ static List<List<SubscriptionCommitContext>> partitionCommitContexts(
+ final List<SubscriptionCommitContext> commitContexts,
+ final int maxBodySize,
+ final int fixedBodySize)
+ throws IOException {
+ if (commitContexts == null || commitContexts.isEmpty()) {
+ return Collections.singletonList(Collections.emptyList());
+ }
+
+ final List<List<SubscriptionCommitContext>> batches = new ArrayList<>();
+ List<SubscriptionCommitContext> currentBatch = new ArrayList<>();
+ long currentBodySize = fixedBodySize;
+ for (final SubscriptionCommitContext commitContext : commitContexts) {
+ final int contextSize =
SubscriptionCommitContext.serialize(commitContext).remaining();
+ if (!currentBatch.isEmpty() && currentBodySize + contextSize >
maxBodySize) {
+ batches.add(currentBatch);
+ currentBatch = new ArrayList<>();
+ currentBodySize = fixedBodySize;
+ }
+ currentBatch.add(commitContext);
+ currentBodySize += contextSize;
+ }
+ batches.add(currentBatch);
+ return batches;
+ }
+
+ static void mergeTopicProgress(
+ final Map<String, TopicProgress> target, final Map<String,
TopicProgress> source) {
+ source.forEach(
+ (topicName, newProgress) ->
+ target.merge(
+ topicName,
+ newProgress,
+ (oldProgress, ignored) -> {
+ final Map<String, RegionProgress> mergedRegionProgress =
+ new LinkedHashMap<>(oldProgress.getRegionProgress());
+ mergedRegionProgress.putAll(newProgress.getRegionProgress());
+ return new TopicProgress(mergedRegionProgress);
+ }));
+ }
+
+ @FunctionalInterface
+ interface PipeSubscribeRequestSender {
+ TPipeSubscribeResp send(TPipeSubscribeReq req) throws TException;
+ }
+
static final class CommitResult {
private static final CommitResult EMPTY =
new CommitResult(Collections.emptyList(), Collections.emptyMap());
diff --git
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqTest.java
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqTest.java
new file mode 100644
index 00000000000..4a29123b70b
--- /dev/null
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/rpc/subscription/payload/request/PipeSubscribeSliceReqTest.java
@@ -0,0 +1,203 @@
+/*
+ * 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.iotdb.rpc.subscription.payload.request;
+
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
+
+import org.apache.tsfile.utils.Binary;
+import org.apache.tsfile.utils.PublicBAOS;
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Optional;
+
+public class PipeSubscribeSliceReqTest {
+
+ @Test
+ public void testBuildAndReassembleSlicesFromNonZeroPosition() throws
Exception {
+ final byte[] bytes = new byte[] {99, 98, 0, 1, 2, 3, 4, 5, 6, 97};
+ final ByteBuffer body = ByteBuffer.wrap(bytes);
+ body.position(2);
+ body.limit(9);
+ final TPipeSubscribeReq req = request(PipeSubscribeRequestType.POLL, body);
+ final int originalPosition = body.position();
+ final int originalLimit = body.limit();
+
+ final int bodySizeLimit = 3;
+ final int sliceCount = PipeSubscribeSliceReqBuilder.getSliceCount(req,
bodySizeLimit);
+ Assert.assertEquals(3, sliceCount);
+ Assert.assertTrue(PipeSubscribeSliceReqBuilder.shouldSlice(req,
bodySizeLimit));
+
+ final PipeSubscribeSliceReqHandler handler = new
PipeSubscribeSliceReqHandler(10);
+ for (int sliceIndex = 0; sliceIndex < sliceCount; sliceIndex++) {
+ final PipeSubscribeSliceReq thriftSlice =
+ PipeSubscribeSliceReqBuilder.buildSliceReq(
+ req, -1, sliceIndex, sliceCount, bodySizeLimit);
+ final int sliceBodyPosition = thriftSlice.body.position();
+ final PipeSubscribeSliceReq slice =
PipeSubscribeSliceReq.fromTPipeSubscribeReq(thriftSlice);
+
+ Assert.assertEquals(sliceBodyPosition, thriftSlice.body.position());
+ Assert.assertEquals(-1, slice.getOrderId());
+ Assert.assertEquals(sliceIndex, slice.getSliceIndex());
+ Assert.assertEquals(sliceCount, slice.getSliceCount());
+ Assert.assertTrue(handler.receiveSlice(slice));
+ if (sliceIndex < sliceCount - 1) {
+ Assert.assertFalse(handler.makeReqIfComplete().isPresent());
+ }
+ }
+
+ final Optional<TPipeSubscribeReq> reassembled =
handler.makeReqIfComplete();
+ Assert.assertTrue(reassembled.isPresent());
+ Assert.assertEquals(req.getVersion(), reassembled.get().getVersion());
+ Assert.assertEquals(req.getType(), reassembled.get().getType());
+ Assert.assertArrayEquals(
+ new byte[] {0, 1, 2, 3, 4, 5, 6},
remainingBytes(reassembled.get().body));
+ Assert.assertEquals(originalPosition, body.position());
+ Assert.assertEquals(originalLimit, body.limit());
+ }
+
+ @Test
+ public void testRejectOutOfOrderSliceAndAcceptFreshRequest() throws
Exception {
+ final TPipeSubscribeReq req =
+ request(PipeSubscribeRequestType.SEEK, ByteBuffer.wrap(new byte[] {0,
1, 2, 3, 4, 5}));
+ final int sliceCount = PipeSubscribeSliceReqBuilder.getSliceCount(req, 2);
+ final PipeSubscribeSliceReqHandler handler = new
PipeSubscribeSliceReqHandler(10);
+
+ Assert.assertTrue(
+ handler.receiveSlice(PipeSubscribeSliceReqBuilder.buildSliceReq(req,
1, 0, sliceCount, 2)));
+ Assert.assertFalse(
+ handler.receiveSlice(PipeSubscribeSliceReqBuilder.buildSliceReq(req,
1, 2, sliceCount, 2)));
+ Assert.assertFalse(handler.makeReqIfComplete().isPresent());
+
+ for (int sliceIndex = 0; sliceIndex < sliceCount; sliceIndex++) {
+ Assert.assertTrue(
+ handler.receiveSlice(
+ PipeSubscribeSliceReqBuilder.buildSliceReq(req, 2, sliceIndex,
sliceCount, 2)));
+ }
+ Assert.assertArrayEquals(
+ new byte[] {0, 1, 2, 3, 4, 5},
remainingBytes(handler.makeReqIfComplete().get().body));
+ }
+
+ @Test
+ public void testRejectChangedMetadataAndBodySizeViolations() throws
Exception {
+ final PipeSubscribeSliceReqHandler handler = new
PipeSubscribeSliceReqHandler(10);
+
+ Assert.assertTrue(
+ handler.receiveSlice(slice(1, PipeSubscribeRequestType.POLL, 4, 0, 2,
new byte[] {0, 1})));
+ Assert.assertFalse(
+ handler.receiveSlice(slice(2, PipeSubscribeRequestType.POLL, 4, 1, 2,
new byte[] {2, 3})));
+
+ Assert.assertFalse(
+ handler.receiveSlice(slice(1, PipeSubscribeRequestType.POLL, 11, 0, 2,
new byte[] {0})));
+ Assert.assertFalse(
+ handler.receiveSlice(slice(1, PipeSubscribeRequestType.POLL, 2, 0, 2,
new byte[] {0, 1})));
+
+ Assert.assertTrue(
+ handler.receiveSlice(slice(1, PipeSubscribeRequestType.POLL, 3, 0, 2,
new byte[] {0})));
+ Assert.assertFalse(
+ handler.receiveSlice(slice(1, PipeSubscribeRequestType.POLL, 3, 1, 2,
new byte[] {1})));
+ Assert.assertFalse(handler.makeReqIfComplete().isPresent());
+ }
+
+ @Test
+ public void testRejectUnsliceableOriginRequestTypes() {
+ for (final PipeSubscribeRequestType type :
+ new PipeSubscribeRequestType[] {
+ PipeSubscribeRequestType.HANDSHAKE,
+ PipeSubscribeRequestType.CLOSE,
+ PipeSubscribeRequestType.SLICE
+ }) {
+ final TPipeSubscribeReq req = request(type, ByteBuffer.wrap(new byte[]
{0, 1}));
+ Assert.assertFalse(PipeSubscribeSliceReqBuilder.shouldSlice(req, 1));
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> PipeSubscribeSliceReqBuilder.buildSliceReq(req, 1, 0, 2, 1));
+ }
+
+ final PipeSubscribeSliceReqHandler handler = new
PipeSubscribeSliceReqHandler(10);
+ Assert.assertFalse(
+ handler.receiveSlice(slice(1, PipeSubscribeRequestType.SLICE, 2, 0, 2,
new byte[] {0})));
+ }
+
+ @Test
+ public void testRejectMalformedSliceEnvelopeBeforeAllocatingBody() throws
Exception {
+ final TPipeSubscribeReq req =
+ request(PipeSubscribeRequestType.POLL, ByteBuffer.wrap(new byte[] {0,
1}));
+ final PipeSubscribeSliceReq malformedSlice =
+ PipeSubscribeSliceReqBuilder.buildSliceReq(req, 1, 0, 2, 1);
+ malformedSlice.body.putInt(malformedSlice.body.position() + 19,
Integer.MAX_VALUE);
+
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> PipeSubscribeSliceReq.fromTPipeSubscribeReq(malformedSlice));
+
+ final PipeSubscribeSliceReq invalidVersionSlice =
+ PipeSubscribeSliceReqBuilder.buildSliceReq(req, 1, 0, 2, 1);
+ invalidVersionSlice.version = (byte) -1;
+ Assert.assertFalse(new
PipeSubscribeSliceReqHandler(10).receiveSlice(invalidVersionSlice));
+ }
+
+ private static TPipeSubscribeReq request(
+ final PipeSubscribeRequestType type, final ByteBuffer body) {
+ final TPipeSubscribeReq req = new TPipeSubscribeReq();
+ req.version = PipeSubscribeRequestVersion.VERSION_1.getVersion();
+ req.type = type.getType();
+ req.body = body;
+ return req;
+ }
+
+ private static PipeSubscribeSliceReq slice(
+ final int orderId,
+ final PipeSubscribeRequestType originType,
+ final int originBodySize,
+ final int sliceIndex,
+ final int sliceCount,
+ final byte[] sliceBody) {
+ final TPipeSubscribeReq thriftReq = new TPipeSubscribeReq();
+ thriftReq.version = PipeSubscribeRequestVersion.VERSION_1.getVersion();
+ thriftReq.type = PipeSubscribeRequestType.SLICE.getType();
+ try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+ final DataOutputStream outputStream = new
DataOutputStream(byteArrayOutputStream)) {
+ ReadWriteIOUtils.write(orderId, outputStream);
+
ReadWriteIOUtils.write(PipeSubscribeRequestVersion.VERSION_1.getVersion(),
outputStream);
+ ReadWriteIOUtils.write(originType.getType(), outputStream);
+ ReadWriteIOUtils.write(originBodySize, outputStream);
+ ReadWriteIOUtils.write(sliceIndex, outputStream);
+ ReadWriteIOUtils.write(sliceCount, outputStream);
+ ReadWriteIOUtils.write(new Binary(sliceBody), outputStream);
+ thriftReq.body =
+ ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0,
byteArrayOutputStream.size());
+ } catch (final IOException e) {
+ throw new AssertionError(e);
+ }
+ return PipeSubscribeSliceReq.fromTPipeSubscribeReq(thriftReq);
+ }
+
+ private static byte[] remainingBytes(final ByteBuffer buffer) {
+ final ByteBuffer duplicate = buffer.duplicate();
+ final byte[] bytes = new byte[duplicate.remaining()];
+ duplicate.get(bytes);
+ return bytes;
+ }
+}
diff --git
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderSlicingTest.java
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderSlicingTest.java
new file mode 100644
index 00000000000..31db29408bc
--- /dev/null
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderSlicingTest.java
@@ -0,0 +1,232 @@
+/*
+ * 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.iotdb.session.subscription.consumer.base;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress;
+import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
+import org.apache.iotdb.rpc.subscription.payload.poll.TopicProgress;
+import org.apache.iotdb.rpc.subscription.payload.poll.WriterId;
+import org.apache.iotdb.rpc.subscription.payload.poll.WriterProgress;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestType;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestVersion;
+import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeSliceReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.SubscriptionHeartbeatReq;
+import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseType;
+import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseVersion;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp;
+
+import org.apache.thrift.TException;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+public class SubscriptionProviderSlicingTest {
+
+ @Test
+ public void testSendSlicesInOrderAndReturnFinalResponse() throws Exception {
+ final TPipeSubscribeReq req = request(new byte[] {0, 1, 2, 3, 4, 5, 6});
+ final List<TPipeSubscribeReq> sentRequests = new ArrayList<>();
+ final TPipeSubscribeResp finalResp = response(TSStatusCode.SUCCESS_STATUS);
+
+ final TPipeSubscribeResp actualResp =
+ AbstractSubscriptionProvider.pipeSubscribeWithOptionalSlicing(
+ req,
+ 3,
+ sentReq -> {
+ sentRequests.add(sentReq);
+ return sentRequests.size() == 3 ? finalResp :
response(TSStatusCode.SUCCESS_STATUS);
+ });
+
+ Assert.assertSame(finalResp, actualResp);
+ Assert.assertEquals(3, sentRequests.size());
+ Integer orderId = null;
+ for (int sliceIndex = 0; sliceIndex < sentRequests.size(); sliceIndex++) {
+ final PipeSubscribeSliceReq slice =
+
PipeSubscribeSliceReq.fromTPipeSubscribeReq(sentRequests.get(sliceIndex));
+ if (orderId == null) {
+ orderId = slice.getOrderId();
+ }
+ Assert.assertEquals(orderId.intValue(), slice.getOrderId());
+ Assert.assertEquals(sliceIndex, slice.getSliceIndex());
+ Assert.assertEquals(sentRequests.size(), slice.getSliceCount());
+ }
+ }
+
+ @Test
+ public void testStopSendingAfterIntermediateSliceError() throws Exception {
+ final TPipeSubscribeReq req = request(new byte[] {0, 1, 2, 3, 4});
+ final List<TPipeSubscribeReq> sentRequests = new ArrayList<>();
+ final TPipeSubscribeResp errorResp =
response(TSStatusCode.SUBSCRIPTION_TYPE_ERROR);
+
+ final TPipeSubscribeResp actualResp =
+ AbstractSubscriptionProvider.pipeSubscribeWithOptionalSlicing(
+ req,
+ 2,
+ sentReq -> {
+ sentRequests.add(sentReq);
+ return sentRequests.size() == 2 ? errorResp :
response(TSStatusCode.SUCCESS_STATUS);
+ });
+
+ Assert.assertSame(errorResp, actualResp);
+ Assert.assertEquals(2, sentRequests.size());
+ }
+
+ @Test
+ public void testSendSmallRequestWithoutWrapping() throws Exception {
+ final TPipeSubscribeReq req = request(new byte[] {0, 1, 2});
+ final List<TPipeSubscribeReq> sentRequests = new ArrayList<>();
+ final TPipeSubscribeResp expectedResp =
response(TSStatusCode.SUCCESS_STATUS);
+
+ final TPipeSubscribeResp actualResp =
+ AbstractSubscriptionProvider.pipeSubscribeWithOptionalSlicing(
+ req,
+ 3,
+ sentReq -> {
+ sentRequests.add(sentReq);
+ return expectedResp;
+ });
+
+ Assert.assertSame(expectedResp, actualResp);
+ Assert.assertEquals(Collections.singletonList(req), sentRequests);
+ }
+
+ @Test
+ public void testRejectOversizedUnsliceableRequest() {
+ for (final PipeSubscribeRequestType type :
+ new PipeSubscribeRequestType[] {
+ PipeSubscribeRequestType.HANDSHAKE, PipeSubscribeRequestType.CLOSE
+ }) {
+ final TPipeSubscribeReq req = request(new byte[] {0, 1, 2, 3});
+ req.type = type.getType();
+ Assert.assertThrows(
+ TException.class,
+ () ->
+ AbstractSubscriptionProvider.pipeSubscribeWithOptionalSlicing(
+ req,
+ 3,
+ ignored -> {
+ Assert.fail();
+ return null;
+ }));
+ }
+ Assert.assertThrows(
+ TException.class,
+ () ->
+ AbstractSubscriptionProvider.pipeSubscribeWithOptionalSlicing(
+ request(new byte[] {0}),
+ 0,
+ ignored -> {
+ Assert.fail();
+ return null;
+ }));
+ }
+
+ @Test
+ public void testPartitionCommitContextsBySerializedSize() throws Exception {
+ final SubscriptionCommitContext first = commitContext("a", 1);
+ final SubscriptionCommitContext second = commitContext("bb", 2);
+ final SubscriptionCommitContext third = commitContext("ccc", 3);
+ final int fixedBodySize = Integer.BYTES;
+ final int maxBodySize =
+ fixedBodySize
+ + SubscriptionCommitContext.serialize(first).remaining()
+ + SubscriptionCommitContext.serialize(second).remaining();
+
+ final List<List<SubscriptionCommitContext>> batches =
+ AbstractSubscriptionProvider.partitionCommitContexts(
+ Arrays.asList(first, second, third), maxBodySize, fixedBodySize);
+
+ Assert.assertEquals(2, batches.size());
+ Assert.assertEquals(Arrays.asList(first, second), batches.get(0));
+ Assert.assertEquals(Collections.singletonList(third), batches.get(1));
+ for (final List<SubscriptionCommitContext> batch : batches) {
+ Assert.assertTrue(
+ SubscriptionHeartbeatReq.toThriftReq(batch).body.remaining() <=
maxBodySize);
+ }
+ Assert.assertEquals(
+ Collections.singletonList(Collections.emptyList()),
+ AbstractSubscriptionProvider.partitionCommitContexts(
+ Collections.emptyList(), maxBodySize, fixedBodySize));
+ }
+
+ @Test
+ public void testMergeTopicProgressAcrossBatches() {
+ final RegionProgress firstRegionProgress = regionProgress("1_1", 1, 1);
+ final RegionProgress updatedFirstRegionProgress = regionProgress("1_1", 1,
2);
+ final RegionProgress secondRegionProgress = regionProgress("1_2", 2, 3);
+ final Map<String, RegionProgress> initialRegions = new LinkedHashMap<>();
+ initialRegions.put("1_1", firstRegionProgress);
+ final Map<String, TopicProgress> target = new LinkedHashMap<>();
+ target.put("topic", new TopicProgress(initialRegions));
+
+ final Map<String, RegionProgress> newRegions = new LinkedHashMap<>();
+ newRegions.put("1_1", updatedFirstRegionProgress);
+ newRegions.put("1_2", secondRegionProgress);
+ final Map<String, TopicProgress> source = new LinkedHashMap<>();
+ source.put("topic", new TopicProgress(newRegions));
+ source.put(
+ "another-topic", new TopicProgress(Collections.singletonMap("2_1",
secondRegionProgress)));
+
+ AbstractSubscriptionProvider.mergeTopicProgress(target, source);
+
+ Assert.assertEquals(2, target.size());
+ Assert.assertEquals(
+ updatedFirstRegionProgress,
target.get("topic").getRegionProgress().get("1_1"));
+ Assert.assertEquals(secondRegionProgress,
target.get("topic").getRegionProgress().get("1_2"));
+ Assert.assertEquals(source.get("another-topic"),
target.get("another-topic"));
+ }
+
+ private static TPipeSubscribeReq request(final byte[] body) {
+ final TPipeSubscribeReq req = new TPipeSubscribeReq();
+ req.version = PipeSubscribeRequestVersion.VERSION_1.getVersion();
+ req.type = PipeSubscribeRequestType.POLL.getType();
+ req.body = ByteBuffer.wrap(body);
+ return req;
+ }
+
+ private static TPipeSubscribeResp response(final TSStatusCode statusCode) {
+ return new TPipeSubscribeResp(
+ new TSStatus(statusCode.getStatusCode()),
+ PipeSubscribeResponseVersion.VERSION_1.getVersion(),
+ PipeSubscribeResponseType.ACK.getType());
+ }
+
+ private static SubscriptionCommitContext commitContext(
+ final String topicName, final long commitId) {
+ return new SubscriptionCommitContext(1, 1, topicName, "group", commitId);
+ }
+
+ private static RegionProgress regionProgress(
+ final String regionId, final int nodeId, final long localSequence) {
+ return new RegionProgress(
+ Collections.singletonMap(
+ new WriterId(regionId, nodeId), new WriterProgress(1L,
localSequence)));
+ }
+}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
index 2c2b00eca12..dd897f3a162 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
@@ -67,6 +67,8 @@ import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeHandshakeR
import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribePollReq;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestType;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestVersion;
+import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeSliceReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeSliceReqHandler;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeSubscribeReq;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeUnsubscribeReq;
import
org.apache.iotdb.rpc.subscription.payload.request.SubscriptionHeartbeatReq;
@@ -110,6 +112,8 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
private static final double POLL_PAYLOAD_SIZE_EXCEED_THRESHOLD = 0.9;
private static final long HEARTBEAT_TIMEOUT_MULTIPLIER = 3L;
+ private static final long MAX_REASSEMBLED_REQUEST_BODY_SIZE = 256L * 1024 *
1024;
+ private static final int MAX_REASSEMBLED_REQUEST_FRAME_COUNT = 4;
private static final IClientManager<ConfigRegionId, ConfigNodeClient>
CONFIG_NODE_CLIENT_MANAGER =
ConfigNodeClientManager.getInstance();
@@ -131,6 +135,18 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
PipeSubscribeResponseVersion.VERSION_1.getVersion(),
PipeSubscribeResponseType.ACK.getType());
+ private static final TPipeSubscribeResp SUBSCRIPTION_SLICE_ACK_RESP =
+ new TPipeSubscribeResp(
+ new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()),
+ PipeSubscribeResponseVersion.VERSION_1.getVersion(),
+ PipeSubscribeResponseType.ACK.getType());
+
+ private static final TPipeSubscribeResp SUBSCRIPTION_SLICE_ERROR_RESP =
+ new TPipeSubscribeResp(
+ new TSStatus(TSStatusCode.SUBSCRIPTION_TYPE_ERROR.getStatusCode()),
+ PipeSubscribeResponseVersion.VERSION_1.getVersion(),
+ PipeSubscribeResponseType.ACK.getType());
+
private final ThreadLocal<ConsumerConfig> consumerConfigThreadLocal = new
ThreadLocal<>();
private final ThreadLocal<PollTimer> pollTimerThreadLocal = new
ThreadLocal<>();
private volatile String authenticatedUsername;
@@ -139,6 +155,8 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
private volatile boolean consumerFenced;
private volatile long lastActivityTimeMs = System.currentTimeMillis();
private final AtomicLong inFlightRequestCount = new AtomicLong(0);
+ private final PipeSubscribeSliceReqHandler sliceReqHandler =
+ new PipeSubscribeSliceReqHandler(getMaxReassembledRequestBodySize());
private long consumerStateVersion;
private static final String SQL_DIALECT_TABLE_VALUE = "table";
@@ -146,6 +164,9 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
@Override
public final TPipeSubscribeResp handle(final TPipeSubscribeReq req) {
final short reqType = req.getType();
+ if (PipeSubscribeRequestType.SLICE.getType() != reqType) {
+ sliceReqHandler.clear();
+ }
final boolean isFencedRequest = beforeHandle(reqType);
try {
if (isFencedRequest) {
@@ -172,6 +193,8 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
return
handlePipeSubscribeClose(PipeSubscribeCloseReq.fromTPipeSubscribeReq(req));
case SEEK:
return
handleSubscriptionSeek(SubscriptionSeekReq.fromThriftReq(req));
+ case SLICE:
+ return handlePipeSubscribeSlice(req);
default:
break;
}
@@ -228,6 +251,7 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
public void invalidateConsumer() {
synchronized (this) {
consumerFenced = true;
+ sliceReqHandler.clear();
clearSharedConsumerState();
}
}
@@ -264,6 +288,7 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
}
}
} finally {
+ sliceReqHandler.clear();
consumerConfigThreadLocal.remove();
pollTimerThreadLocal.remove();
authenticatedUsername = null;
@@ -271,6 +296,33 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
}
}
+ private TPipeSubscribeResp handlePipeSubscribeSlice(final TPipeSubscribeReq
thriftReq) {
+ final PipeSubscribeSliceReq sliceReq;
+ try {
+ sliceReq = PipeSubscribeSliceReq.fromTPipeSubscribeReq(thriftReq);
+ } catch (final RuntimeException e) {
+ sliceReqHandler.clear();
+ return SUBSCRIPTION_SLICE_ERROR_RESP;
+ }
+ if (!sliceReqHandler.receiveSlice(sliceReq)) {
+ return SUBSCRIPTION_SLICE_ERROR_RESP;
+ }
+ return sliceReqHandler
+ .makeReqIfComplete()
+ .map(this::handle)
+ .orElse(SUBSCRIPTION_SLICE_ACK_RESP);
+ }
+
+ private static int getMaxReassembledRequestBodySize() {
+ final long maxFrameSize =
IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize();
+ final long desiredSize =
+ Math.min(
+ MAX_REASSEMBLED_REQUEST_BODY_SIZE, maxFrameSize *
MAX_REASSEMBLED_REQUEST_FRAME_COUNT);
+ final long memoryLimit = Math.max(maxFrameSize,
Runtime.getRuntime().maxMemory() / 16);
+ return (int)
+ Math.min(Integer.MAX_VALUE, Math.max(maxFrameSize,
Math.min(desiredSize, memoryLimit)));
+ }
+
@Override
public void handleTimeout() {
synchronized (this) {
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
index 8b2bd022241..a8c7dc7841f 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
@@ -30,7 +30,10 @@ import
org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
import org.apache.iotdb.rpc.subscription.config.TopicConstant;
import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeHandshakeReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeSliceReqBuilder;
import
org.apache.iotdb.rpc.subscription.payload.request.SubscriptionHeartbeatReq;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp;
import org.junit.Assert;
import org.junit.Test;
@@ -221,6 +224,58 @@ public class SubscriptionReceiverV1Test {
receiver.handle(SubscriptionHeartbeatReq.toThriftReq()).getStatus().getCode());
}
+ @Test
+ public void testHandleSlicesDispatchesReassembledRequest() throws Exception {
+ final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final TPipeSubscribeReq heartbeatReq =
+ SubscriptionHeartbeatReq.toThriftReq(Collections.emptyList());
+ final int bodySizeLimit = 2;
+ final int sliceCount =
PipeSubscribeSliceReqBuilder.getSliceCount(heartbeatReq, bodySizeLimit);
+
+ final TPipeSubscribeResp firstResp =
+ receiver.handle(
+ PipeSubscribeSliceReqBuilder.buildSliceReq(
+ heartbeatReq, 1, 0, sliceCount, bodySizeLimit));
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
firstResp.getStatus().getCode());
+
+ final TPipeSubscribeResp lastResp =
+ receiver.handle(
+ PipeSubscribeSliceReqBuilder.buildSliceReq(
+ heartbeatReq, 1, 1, sliceCount, bodySizeLimit));
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_MISSING_CONSUMER.getStatusCode(),
lastResp.getStatus().getCode());
+ }
+
+ @Test
+ public void testNonSliceRequestClearsIncompleteSlices() throws Exception {
+ final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final TPipeSubscribeReq heartbeatReq =
+ SubscriptionHeartbeatReq.toThriftReq(Collections.emptyList());
+ final int bodySizeLimit = 2;
+ final int sliceCount =
PipeSubscribeSliceReqBuilder.getSliceCount(heartbeatReq, bodySizeLimit);
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ receiver
+ .handle(
+ PipeSubscribeSliceReqBuilder.buildSliceReq(
+ heartbeatReq, 1, 0, sliceCount, bodySizeLimit))
+ .getStatus()
+ .getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_MISSING_CONSUMER.getStatusCode(),
+
receiver.handle(SubscriptionHeartbeatReq.toThriftReq()).getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_TYPE_ERROR.getStatusCode(),
+ receiver
+ .handle(
+ PipeSubscribeSliceReqBuilder.buildSliceReq(
+ heartbeatReq, 1, 1, sliceCount, bodySizeLimit))
+ .getStatus()
+ .getCode());
+ }
+
@Test
public void testCalculateConsumerInactivityTimeoutUsesDefaultTimeout()
throws Exception {
final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
