Skip to content

Commit cebe663

Browse files
CaideyipiCaideyipi
authored
Fix hybrid pipe queue memory admission across data regions (#19080)
1 parent 6466747 commit cebe663

2 files changed

Lines changed: 204 additions & 1 deletion

File tree

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgent.java‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1232,6 +1232,16 @@ private void calculateInsertNodeQueueMemory(
12321232
return;
12331233
}
12341234

1235+
// Snapshot/query and non-streaming sources do not retain InsertNode tablets.
1236+
if (isSnapshotMode(sourceParameters)
1237+
|| !sourceParameters.getBooleanOrDefault(
1238+
Arrays.asList(
1239+
PipeSourceConstant.EXTRACTOR_MODE_STREAMING_KEY,
1240+
PipeSourceConstant.SOURCE_MODE_STREAMING_KEY),
1241+
PipeSourceConstant.EXTRACTOR_MODE_STREAMING_DEFAULT_VALUE)) {
1242+
return;
1243+
}
1244+
12351245
// If the realtime mode is batch or file, we do not need to allocate memory
12361246
final String realtimeMode =
12371247
sourceParameters.getStringByKeys(
@@ -1242,8 +1252,15 @@ private void calculateInsertNodeQueueMemory(
12421252
return;
12431253
}
12441254

1255+
// Hybrid sources share a pipe-wide floating-memory limit and downgrade to TsFiles when it is
1256+
// reached. Charging a full queue for every region rejects an idle pipe as its region count
1257+
// grows.
1258+
// Only forced-log sources cannot downgrade and retain the per-region estimate.
12451259
final long needFloatingMemory =
1246-
PipeConfig.getInstance().getPipeInsertNodeQueueMemory() * dataRegionTaskCount;
1260+
PipeConfig.getInstance().getPipeInsertNodeQueueMemory()
1261+
* (PipeSourceConstant.EXTRACTOR_REALTIME_MODE_FORCED_LOG_VALUE.equals(realtimeMode)
1262+
? dataRegionTaskCount
1263+
: 1);
12471264
final long allocatedMemorySizeInBytes = this.getAllFloatingMemoryUsageInByte();
12481265
final long remainingMemory =
12491266
PipeDataNodeResourceManager.memory().getTotalFloatingMemorySizeInBytes()

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeDataNodeTaskAgentMemoryTest.java‎

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
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.db.pipe.agent.task;
21+
22+
import org.apache.iotdb.commons.conf.CommonConfig;
23+
import org.apache.iotdb.commons.conf.CommonDescriptor;
24+
import org.apache.iotdb.commons.pipe.config.constant.PipeSourceConstant;
25+
import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
26+
import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
27+
import org.apache.iotdb.pipe.api.exception.PipeException;
28+
29+
import org.junit.After;
30+
import org.junit.Assert;
31+
import org.junit.Before;
32+
import org.junit.Test;
33+
34+
import java.lang.reflect.InvocationTargetException;
35+
import java.lang.reflect.Method;
36+
import java.util.Collections;
37+
import java.util.Map;
38+
39+
public class PipeDataNodeTaskAgentMemoryTest {
40+
41+
private static final int DATA_REGION_COUNT = 14;
42+
43+
private final CommonConfig config = CommonDescriptor.getInstance().getConfig();
44+
private final PipeDataNodeTaskAgent taskAgent = new PipeDataNodeTaskAgent();
45+
private long originalQueueMemory;
46+
private double originalFloatingMemoryProportion;
47+
private long floatingMemorySizeInBytes;
48+
private Method memoryCheck;
49+
50+
@Before
51+
public void setUp() throws Exception {
52+
originalQueueMemory = config.getPipeInsertNodeQueueMemory();
53+
originalFloatingMemoryProportion = config.getPipeTotalFloatingMemoryProportion();
54+
config.setPipeTotalFloatingMemoryProportion(0.5);
55+
floatingMemorySizeInBytes =
56+
PipeDataNodeResourceManager.memory().getTotalFloatingMemorySizeInBytes();
57+
Assert.assertTrue(floatingMemorySizeInBytes > 1);
58+
config.setPipeInsertNodeQueueMemory(floatingMemorySizeInBytes / 2);
59+
memoryCheck =
60+
PipeDataNodeTaskAgent.class.getDeclaredMethod(
61+
"calculateInsertNodeQueueMemory", PipeParameters.class, int.class);
62+
memoryCheck.setAccessible(true);
63+
}
64+
65+
@After
66+
public void tearDown() {
67+
config.setPipeInsertNodeQueueMemory(originalQueueMemory);
68+
config.setPipeTotalFloatingMemoryProportion(originalFloatingMemoryProportion);
69+
}
70+
71+
@Test
72+
public void testHybridQueueBudgetIsSharedAcrossDataRegions() throws Exception {
73+
checkMemory(Collections.emptyMap(), DATA_REGION_COUNT);
74+
for (final String key :
75+
new String[] {
76+
PipeSourceConstant.SOURCE_MODE_STREAMING_KEY,
77+
PipeSourceConstant.EXTRACTOR_MODE_STREAMING_KEY
78+
}) {
79+
checkMemory(Collections.singletonMap(key, "true"), DATA_REGION_COUNT);
80+
}
81+
for (final String key :
82+
new String[] {
83+
PipeSourceConstant.SOURCE_REALTIME_MODE_KEY,
84+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_KEY
85+
}) {
86+
for (final String mode :
87+
new String[] {
88+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_HYBRID_VALUE,
89+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_STREAM_MODE_VALUE,
90+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_LOG_VALUE
91+
}) {
92+
checkMemory(Collections.singletonMap(key, mode), DATA_REGION_COUNT);
93+
}
94+
}
95+
}
96+
97+
@Test
98+
public void testHybridQueueStillRequiresAvailableFloatingMemory() {
99+
config.setPipeInsertNodeQueueMemory(floatingMemorySizeInBytes + 1);
100+
final PipeException exception =
101+
Assert.assertThrows(
102+
PipeException.class, () -> checkMemory(Collections.emptyMap(), DATA_REGION_COUNT));
103+
Assert.assertTrue(
104+
exception.getMessage().contains(Long.toString(floatingMemorySizeInBytes + 1)));
105+
}
106+
107+
@Test
108+
public void testForcedLogStillRequiresPerRegionMemory() throws Exception {
109+
for (final String key :
110+
new String[] {
111+
PipeSourceConstant.SOURCE_REALTIME_MODE_KEY,
112+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_KEY
113+
}) {
114+
final Map<String, String> attributes =
115+
Collections.singletonMap(
116+
key, PipeSourceConstant.EXTRACTOR_REALTIME_MODE_FORCED_LOG_VALUE);
117+
checkMemory(attributes, 1);
118+
final PipeException exception =
119+
Assert.assertThrows(
120+
PipeException.class, () -> checkMemory(attributes, DATA_REGION_COUNT));
121+
Assert.assertTrue(
122+
exception
123+
.getMessage()
124+
.contains(Long.toString(config.getPipeInsertNodeQueueMemory() * DATA_REGION_COUNT)));
125+
}
126+
}
127+
128+
@Test
129+
public void testSourcesWithoutTabletQueuesSkipFloatingMemoryCheck() throws Exception {
130+
config.setPipeInsertNodeQueueMemory(floatingMemorySizeInBytes + 1);
131+
checkMemory(Collections.emptyMap(), 0);
132+
for (final String key :
133+
new String[] {
134+
PipeSourceConstant.SOURCE_REALTIME_ENABLE_KEY,
135+
PipeSourceConstant.EXTRACTOR_REALTIME_ENABLE_KEY
136+
}) {
137+
checkMemory(Collections.singletonMap(key, "false"), DATA_REGION_COUNT);
138+
}
139+
for (final String mode :
140+
new String[] {
141+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_FILE_VALUE,
142+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_BATCH_MODE_VALUE
143+
}) {
144+
for (final String key :
145+
new String[] {
146+
PipeSourceConstant.SOURCE_REALTIME_MODE_KEY,
147+
PipeSourceConstant.EXTRACTOR_REALTIME_MODE_KEY
148+
}) {
149+
checkMemory(Collections.singletonMap(key, mode), DATA_REGION_COUNT);
150+
}
151+
}
152+
for (final String key :
153+
new String[] {
154+
PipeSourceConstant.SOURCE_MODE_STREAMING_KEY,
155+
PipeSourceConstant.EXTRACTOR_MODE_STREAMING_KEY
156+
}) {
157+
checkMemory(Collections.singletonMap(key, "false"), DATA_REGION_COUNT);
158+
}
159+
for (final String key :
160+
new String[] {
161+
PipeSourceConstant.SOURCE_MODE_SNAPSHOT_KEY,
162+
PipeSourceConstant.EXTRACTOR_MODE_SNAPSHOT_KEY
163+
}) {
164+
checkMemory(Collections.singletonMap(key, "true"), DATA_REGION_COUNT);
165+
}
166+
for (final String key :
167+
new String[] {PipeSourceConstant.SOURCE_MODE_KEY, PipeSourceConstant.EXTRACTOR_MODE_KEY}) {
168+
for (final String mode :
169+
new String[] {
170+
PipeSourceConstant.EXTRACTOR_MODE_SNAPSHOT_VALUE,
171+
PipeSourceConstant.EXTRACTOR_MODE_QUERY_VALUE
172+
}) {
173+
checkMemory(Collections.singletonMap(key, mode), DATA_REGION_COUNT);
174+
}
175+
}
176+
}
177+
178+
private void checkMemory(final Map<String, String> attributes, final int dataRegionCount)
179+
throws Exception {
180+
try {
181+
memoryCheck.invoke(taskAgent, new PipeParameters(attributes), dataRegionCount);
182+
} catch (final InvocationTargetException e) {
183+
throw (Exception) e.getCause();
184+
}
185+
}
186+
}

0 commit comments

Comments
 (0)