Skip to content

Commit 566ea99

Browse files
fix
1 parent becf7f3 commit 566ea99

2 files changed

Lines changed: 217 additions & 0 deletions

File tree

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.rpc.subscription.payload.request;
21+
22+
import org.apache.iotdb.rpc.subscription.payload.poll.TopicProgress;
23+
import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
24+
25+
import org.apache.tsfile.utils.PublicBAOS;
26+
import org.apache.tsfile.utils.ReadWriteIOUtils;
27+
28+
import java.io.DataOutputStream;
29+
import java.io.IOException;
30+
import java.nio.ByteBuffer;
31+
import java.util.Collections;
32+
import java.util.Objects;
33+
34+
public class SubscriptionSeekReq extends TPipeSubscribeReq {
35+
36+
/** Seek type constants. */
37+
public static final short SEEK_TO_BEGINNING = 1;
38+
39+
public static final short SEEK_TO_END = 2;
40+
public static final short SEEK_TO_TIMESTAMP = 3;
41+
public static final short SEEK_TO_TOPIC_PROGRESS = 4;
42+
public static final short SEEK_AFTER_TOPIC_PROGRESS = 5;
43+
44+
private transient String topicName;
45+
private transient short seekType;
46+
private transient long timestamp; // only meaningful when seekType == SEEK_TO_TIMESTAMP
47+
private transient TopicProgress topicProgress = new TopicProgress(Collections.emptyMap());
48+
49+
public String getTopicName() {
50+
return topicName;
51+
}
52+
53+
public short getSeekType() {
54+
return seekType;
55+
}
56+
57+
public long getTimestamp() {
58+
return timestamp;
59+
}
60+
61+
public TopicProgress getTopicProgress() {
62+
return topicProgress;
63+
}
64+
65+
/////////////////////////////// Thrift ///////////////////////////////
66+
67+
/**
68+
* Serialize the incoming parameters into {@code TPipeSubscribeReq}, called by the subscription
69+
* client.
70+
*/
71+
public static SubscriptionSeekReq toThriftReq(
72+
final String topicName, final short seekType, final long timestamp) throws IOException {
73+
return toThriftReq(topicName, seekType, timestamp, null);
74+
}
75+
76+
public static SubscriptionSeekReq toTopicProgressReq(
77+
final String topicName, final TopicProgress topicProgress) throws IOException {
78+
return toThriftReq(topicName, SEEK_TO_TOPIC_PROGRESS, 0, topicProgress);
79+
}
80+
81+
public static SubscriptionSeekReq toSeekAfterTopicProgressReq(
82+
final String topicName, final TopicProgress topicProgress) throws IOException {
83+
return toThriftReq(topicName, SEEK_AFTER_TOPIC_PROGRESS, 0, topicProgress);
84+
}
85+
86+
public static SubscriptionSeekReq toThriftReq(
87+
final String topicName,
88+
final short seekType,
89+
final long timestamp,
90+
final TopicProgress topicProgress)
91+
throws IOException {
92+
final SubscriptionSeekReq req = new SubscriptionSeekReq();
93+
94+
req.topicName = topicName;
95+
req.seekType = seekType;
96+
req.timestamp = timestamp;
97+
req.topicProgress =
98+
Objects.nonNull(topicProgress) ? topicProgress : new TopicProgress(Collections.emptyMap());
99+
100+
req.version = PipeSubscribeRequestVersion.VERSION_1.getVersion();
101+
req.type = PipeSubscribeRequestType.SEEK.getType();
102+
try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
103+
final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
104+
ReadWriteIOUtils.write(topicName, outputStream);
105+
ReadWriteIOUtils.write(seekType, outputStream);
106+
if (seekType == SEEK_TO_TIMESTAMP) {
107+
ReadWriteIOUtils.write(timestamp, outputStream);
108+
} else if (seekType == SEEK_TO_TOPIC_PROGRESS || seekType == SEEK_AFTER_TOPIC_PROGRESS) {
109+
req.topicProgress.serialize(outputStream);
110+
}
111+
req.body = ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size());
112+
}
113+
114+
return req;
115+
}
116+
117+
/**
118+
* Deserialize {@code TPipeSubscribeReq} to obtain parameters, called by the subscription server.
119+
*/
120+
public static SubscriptionSeekReq fromThriftReq(final TPipeSubscribeReq seekReq) {
121+
final SubscriptionSeekReq req = new SubscriptionSeekReq();
122+
123+
if (Objects.nonNull(seekReq.body) && seekReq.body.hasRemaining()) {
124+
req.topicName = ReadWriteIOUtils.readString(seekReq.body);
125+
req.seekType = ReadWriteIOUtils.readShort(seekReq.body);
126+
if (req.seekType == SEEK_TO_TIMESTAMP) {
127+
req.timestamp = ReadWriteIOUtils.readLong(seekReq.body);
128+
} else if (req.seekType == SEEK_TO_TOPIC_PROGRESS
129+
|| req.seekType == SEEK_AFTER_TOPIC_PROGRESS) {
130+
req.topicProgress = TopicProgress.deserialize(seekReq.body);
131+
}
132+
}
133+
134+
req.version = seekReq.version;
135+
req.type = seekReq.type;
136+
137+
return req;
138+
}
139+
140+
/////////////////////////////// Object ///////////////////////////////
141+
142+
@Override
143+
public boolean equals(final Object obj) {
144+
if (this == obj) {
145+
return true;
146+
}
147+
if (obj == null || getClass() != obj.getClass()) {
148+
return false;
149+
}
150+
final SubscriptionSeekReq that = (SubscriptionSeekReq) obj;
151+
return Objects.equals(this.topicName, that.topicName)
152+
&& this.seekType == that.seekType
153+
&& this.timestamp == that.timestamp
154+
&& Objects.equals(this.topicProgress, that.topicProgress)
155+
&& this.version == that.version
156+
&& this.type == that.type
157+
&& Objects.equals(this.body, that.body);
158+
}
159+
160+
@Override
161+
public int hashCode() {
162+
return Objects.hash(topicName, seekType, timestamp, topicProgress, version, type, body);
163+
}
164+
}
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.rpc.subscription.payload.request;
21+
22+
import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress;
23+
import org.apache.iotdb.rpc.subscription.payload.poll.TopicProgress;
24+
import org.apache.iotdb.rpc.subscription.payload.poll.WriterId;
25+
import org.apache.iotdb.rpc.subscription.payload.poll.WriterProgress;
26+
27+
import org.junit.Test;
28+
29+
import java.io.IOException;
30+
import java.util.Collections;
31+
import java.util.LinkedHashMap;
32+
import java.util.Map;
33+
34+
import static org.junit.Assert.assertEquals;
35+
36+
public class SubscriptionSeekReqTest {
37+
38+
@Test
39+
public void testTopicProgressSeekRoundTrip() throws IOException {
40+
final Map<WriterId, WriterProgress> writerPositions = new LinkedHashMap<>();
41+
writerPositions.put(new WriterId("1_100", 1), new WriterProgress(1000L, 10L));
42+
final TopicProgress original =
43+
new TopicProgress(Collections.singletonMap("1_100", new RegionProgress(writerPositions)));
44+
45+
final SubscriptionSeekReq req =
46+
SubscriptionSeekReq.toSeekAfterTopicProgressReq("topicA", original);
47+
final SubscriptionSeekReq parsed = SubscriptionSeekReq.fromThriftReq(req);
48+
49+
assertEquals(SubscriptionSeekReq.SEEK_AFTER_TOPIC_PROGRESS, parsed.getSeekType());
50+
assertEquals("topicA", parsed.getTopicName());
51+
assertEquals(original, parsed.getTopicProgress());
52+
}
53+
}

0 commit comments

Comments
 (0)