Skip to content

Fix oversized Load TsFile piece dispatch #18627

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Caideyipi wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix oversized Load TsFile piece dispatch #18627

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion ...rc/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -629,10 +629,30 @@ public TSchemaFetchResponse fetchSchema(final TSchemaFetchRequest req) {

@Override
public TLoadResp sendTsFilePieceNode(final TTsFilePieceReq req) {
LOGGER.info(DataNodeMiscMessages.RECEIVE_LOAD_NODE, req.uuid);
if (!req.isSetSliceIndex() || req.sliceIndex == 0) {
LOGGER.info(DataNodeMiscMessages.RECEIVE_LOAD_NODE, req.uuid);
}

final ConsensusGroupId groupId =
ConsensusGroupId.Factory.createFromTConsensusGroupId(req.consensusGroupId);
final boolean isSliced =
req.isSetSliceIndex() || req.isSetSliceCount() || req.isSetOriginBodySize();
if (isSliced) {
if (!req.isSetSliceIndex() || !req.isSetSliceCount() || !req.isSetOriginBodySize()) {
return createTLoadResp(
new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode()));
}
return createTLoadResp(
StorageEngine.getInstance()
.writeLoadTsFileNodeSlice(
(DataRegionId) groupId,
req.body,
req.uuid,
req.sliceIndex,
req.sliceCount,
req.originBodySize));
}

final LoadTsFilePieceNode pieceNode = (LoadTsFilePieceNode) PlanNodeType.deserialize(req.body);
if (pieceNode == null) {
return createTLoadResp(
Expand Down
102 changes: 82 additions & 20 deletions ...in/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

package org.apache.iotdb.db.queryengine.plan.scheduler.load;

import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
Expand Down Expand Up @@ -63,6 +64,7 @@
import java.io.IOException;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
Expand All @@ -82,6 +84,7 @@ public class LoadTsFileDispatcherImpl implements IFragInstanceDispatcher, AutoCl

private static final int MAX_CONNECTION_TIMEOUT_MS = 24 * 60 * 60 * 1000; // 1 day
private static final int FIRST_ADJUSTMENT_TIMEOUT_MS = 6 * 60 * 60 * 1000; // 6 hours
private static final int LOAD_TSFILE_PIECE_RPC_FRAME_RESERVED_BYTES = 1024;
private static final AtomicInteger CONNECTION_TIMEOUT_MS =
new AtomicInteger(IoTDBDescriptor.getInstance().getConfig().getConnectionTimeoutInMS());

Expand Down Expand Up @@ -143,26 +146,83 @@ public Future<FragInstanceDispatchResult> dispatch(

private void dispatchOneInstance(FragmentInstance instance)
throws FragmentInstanceDispatchException {
TTsFilePieceReq loadTsFileReq = null;
List<TTsFilePieceReq> loadTsFileReqs = null;

for (TDataNodeLocation dataNodeLocation :
instance.getRegionReplicaSet().getDataNodeLocations()) {
TEndPoint endPoint = dataNodeLocation.getInternalEndPoint();
if (isDispatchedToLocal(endPoint)) {
dispatchLocally(instance);
} else {
if (loadTsFileReq == null) {
loadTsFileReq =
new TTsFilePieceReq(
if (loadTsFileReqs == null) {
loadTsFileReqs =
splitTsFilePieceReq(
instance.getFragment().getPlanNodeTree().serializeToByteBuffer(),
uuid,
instance.getRegionReplicaSet().getRegionId());
instance.getRegionReplicaSet().getRegionId(),
getLoadTsFilePieceBodySizeLimit());
}
dispatchRemote(loadTsFileReq, endPoint);
dispatchRemote(loadTsFileReqs, endPoint);
}
}
}

private static int getLoadTsFilePieceBodySizeLimit() {
final int thriftMaxFrameSize =
IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize();
return Math.max(1, thriftMaxFrameSize - LOAD_TSFILE_PIECE_RPC_FRAME_RESERVED_BYTES);
}

static List<TTsFilePieceReq> splitTsFilePieceReq(
final ByteBuffer body,
final String uuid,
final TConsensusGroupId consensusGroupId,
final int bodySizeLimit) {
if (bodySizeLimit <= 0) {
throw new IllegalArgumentException();
}

final int originBodySize = body.remaining();
final int sliceCount = getSliceCount(originBodySize, bodySizeLimit);
final List<TTsFilePieceReq> requests = new ArrayList<>(sliceCount);
if (sliceCount == 1) {
requests.add(createTsFilePieceReq(body.duplicate(), uuid, consensusGroupId));
return requests;
}

final int originPosition = body.position();
for (int sliceIndex = 0; sliceIndex < sliceCount; sliceIndex++) {
final int startOffset = sliceIndex * bodySizeLimit;
final int endOffset = startOffset + Math.min(bodySizeLimit, originBodySize - startOffset);
final ByteBuffer slicedBody = body.duplicate();
slicedBody.position(originPosition + startOffset);
slicedBody.limit(originPosition + endOffset);
requests.add(
createTsFilePieceReq(slicedBody.slice(), uuid, consensusGroupId)
.setSliceIndex(sliceIndex)
.setSliceCount(sliceCount)
.setOriginBodySize(originBodySize));
}
return requests;
}

static int getSliceCount(final int bodySize, final int bodySizeLimit) {
if (bodySize < 0 || bodySizeLimit <= 0) {
throw new IllegalArgumentException();
}
return bodySize == 0 ? 1 : (bodySize - 1) / bodySizeLimit + 1;
}

private static TTsFilePieceReq createTsFilePieceReq(
final ByteBuffer body, final String uuid, final TConsensusGroupId consensusGroupId) {
final TTsFilePieceReq request =
new TTsFilePieceReq().setUuid(uuid).setConsensusGroupId(consensusGroupId);
// The generated setter copies the whole buffer, while these immutable slices remain valid until
// all replicas have been dispatched.
request.body = body;
return request;
}

public void dispatchLocally(FragmentInstance instance) throws FragmentInstanceDispatchException {
if (isGeneratedByPipe) {
LOGGER.debug(DataNodeQueryMessages.RECEIVE_LOAD_NODE_FROM_UUID, uuid);
Expand Down Expand Up @@ -222,25 +282,27 @@ public void dispatchLocally(FragmentInstance instance) throws FragmentInstanceDi
}
}

private void dispatchRemote(TTsFilePieceReq loadTsFileReq, TEndPoint endPoint)
private void dispatchRemote(List<TTsFilePieceReq> loadTsFileReqs, TEndPoint endPoint)
throws FragmentInstanceDispatchException {
boolean transferAttemptRecorded = false;
try (SyncDataNodeInternalServiceClient client =
internalServiceClientManager.borrowClient(endPoint)) {
client.setTimeout(CONNECTION_TIMEOUT_MS.get());

final TLoadResp loadResp = client.sendTsFilePieceNode(loadTsFileReq);
if (!loadResp.isAccepted()) {
recordTransferAttempt(
endPoint,
false,
loadResp.isSetStatus()
? String.valueOf(loadResp.getStatus().getCode())
: UserDataTransferErrorCode.REMOTE_REJECTED.name(),
null);
transferAttemptRecorded = true;
LOGGER.warn(loadResp.message);
throw new FragmentInstanceDispatchException(loadResp.status);
for (final TTsFilePieceReq loadTsFileReq : loadTsFileReqs) {
final TLoadResp loadResp = client.sendTsFilePieceNode(loadTsFileReq);
if (!loadResp.isAccepted()) {
recordTransferAttempt(
endPoint,
false,
loadResp.isSetStatus()
? String.valueOf(loadResp.getStatus().getCode())
: UserDataTransferErrorCode.REMOTE_REJECTED.name(),
null);
transferAttemptRecorded = true;
LOGGER.warn(loadResp.message);
throw new FragmentInstanceDispatchException(loadResp.status);
}
}
recordTransferAttempt(endPoint, true, null, null);
transferAttemptRecorded = true;
Expand All @@ -253,7 +315,7 @@ private void dispatchRemote(TTsFilePieceReq loadTsFileReq, TEndPoint endPoint)
final String exceptionMessage =
String.format(
"failed to dispatch load command %s to node %s because of exception: %s",
loadTsFileReq, endPoint, e);
uuid, endPoint, e);
LOGGER.warn(exceptionMessage, e);
throw new FragmentInstanceDispatchException(
new TSStatus()
Expand Down
31 changes: 31 additions & 0 deletions iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import org.apache.iotdb.commons.exception.ShutdownException;
import org.apache.iotdb.commons.exception.StartupException;
import org.apache.iotdb.commons.file.SystemFileFactory;
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType;
import org.apache.iotdb.commons.schema.ttl.TTLCache;
import org.apache.iotdb.commons.service.IService;
import org.apache.iotdb.commons.service.ServiceType;
Expand Down Expand Up @@ -79,6 +80,7 @@
import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALException;
import org.apache.iotdb.db.storageengine.dataregion.wal.recover.WALRecoverManager;
import org.apache.iotdb.db.storageengine.load.LoadTsFileManager;
import org.apache.iotdb.db.storageengine.load.LoadTsFilePieceNodeAssembler;
import org.apache.iotdb.db.storageengine.load.limiter.LoadTsFileRateLimiter;
import org.apache.iotdb.db.storageengine.rescon.disk.TierManager;
import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo;
Expand All @@ -97,6 +99,7 @@
import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
Expand Down Expand Up @@ -1066,6 +1069,34 @@ public TSStatus writeLoadTsFileNode(
return RpcUtils.SUCCESS_STATUS;
}

public TSStatus writeLoadTsFileNodeSlice(
final DataRegionId dataRegionId,
final ByteBuffer body,
final String uuid,
final int sliceIndex,
final int sliceCount,
final int originBodySize) {
final LoadTsFilePieceNodeAssembler.Result result =
loadTsFileManager.appendPieceNodeSlice(
dataRegionId, uuid, body, sliceIndex, sliceCount, originBodySize);
if (!result.isValid()) {
return new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode());
}
if (!result.isComplete()) {
return RpcUtils.SUCCESS_STATUS;
}

try {
final Object planNode = PlanNodeType.deserialize(result.getBody());
if (!(planNode instanceof LoadTsFilePieceNode)) {
return new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode());
}
return writeLoadTsFileNode(dataRegionId, (LoadTsFilePieceNode) planNode, uuid);
} catch (final Exception e) {
return new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode());
}
}

public TSStatus executeLoadCommand(
LoadTsFileScheduler.LoadCommand loadCommand,
String uuid,
Expand Down
61 changes: 59 additions & 2 deletions ...core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
import org.apache.iotdb.commons.conf.IoTDBConstant;
import org.apache.iotdb.commons.consensus.ConsensusGroupId;
import org.apache.iotdb.commons.consensus.DataRegionId;
import org.apache.iotdb.commons.consensus.index.ProgressIndex;
import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
import org.apache.iotdb.commons.disk.FolderManager;
Expand Down Expand Up @@ -77,6 +78,7 @@

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.DirectoryNotEmptyException;
import java.nio.file.Files;
import java.nio.file.Path;
Expand Down Expand Up @@ -123,6 +125,9 @@ public class LoadTsFileManager {

private final Map<String, TsFileWriterManager> uuid2WriterManager = new ConcurrentHashMap<>();

private final Map<String, Map<DataRegionId, LoadTsFilePieceNodeAssembler>>
uuid2PieceNodeAssembler = new ConcurrentHashMap<>();

private final Map<String, CleanupTask> uuid2CleanupTask = new ConcurrentHashMap<>();
private final PriorityBlockingQueue<CleanupTask> cleanupTaskQueue = new PriorityBlockingQueue<>();

Expand All @@ -145,6 +150,7 @@ public void stop() {
cleanupTaskQueue.clear();
}
new HashSet<>(uuid2WriterManager.keySet()).forEach(this::forceCloseWriterManager);
uuid2PieceNodeAssembler.clear();
}

private long getCleanupTaskDelayInMs() {
Expand Down Expand Up @@ -301,6 +307,53 @@ public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNo
}
}

public LoadTsFilePieceNodeAssembler.Result appendPieceNodeSlice(
final DataRegionId dataRegionId,
final String uuid,
final ByteBuffer body,
final int sliceIndex,
final int sliceCount,
final int originBodySize) {
createCleanupTaskIfAbsent(uuid);

final Optional<CleanupTask> cleanupTask = Optional.ofNullable(uuid2CleanupTask.get(uuid));
cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning);
try {
final Map<DataRegionId, LoadTsFilePieceNodeAssembler> regionId2Assembler =
uuid2PieceNodeAssembler.computeIfAbsent(uuid, key -> new ConcurrentHashMap<>());
synchronized (regionId2Assembler) {
final LoadTsFilePieceNodeAssembler assembler;
if (sliceIndex == 0) {
assembler = new LoadTsFilePieceNodeAssembler(sliceCount, originBodySize);
regionId2Assembler.put(dataRegionId, assembler);
} else {
assembler = regionId2Assembler.get(dataRegionId);
if (assembler == null) {
removePieceNodeAssemblerIfEmpty(uuid, regionId2Assembler);
return LoadTsFilePieceNodeAssembler.Result.invalid();
}
}

final LoadTsFilePieceNodeAssembler.Result result =
assembler.append(body, sliceIndex, sliceCount, originBodySize);
if (!result.isValid() || result.isComplete()) {
regionId2Assembler.remove(dataRegionId, assembler);
removePieceNodeAssemblerIfEmpty(uuid, regionId2Assembler);
}
return result;
}
} finally {
cleanupTask.ifPresent(CleanupTask::markLoadTaskNotRunning);
}
}

private void removePieceNodeAssemblerIfEmpty(
final String uuid, final Map<DataRegionId, LoadTsFilePieceNodeAssembler> regionId2Assembler) {
if (regionId2Assembler.isEmpty()) {
uuid2PieceNodeAssembler.remove(uuid, regionId2Assembler);
}
}

private FolderManager getFolderManager() throws DiskSpaceInsufficientException {
if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) {
synchronized (FOLDER_MANAGER) {
Expand Down Expand Up @@ -333,7 +386,7 @@ public boolean loadAll(
boolean isGeneratedByPipe,
Map<TTimePartitionSlot, ProgressIndex> timePartitionProgressIndexMap)
throws IOException, LoadFileException {
if (!uuid2WriterManager.containsKey(uuid)) {
if (!uuid2WriterManager.containsKey(uuid) || uuid2PieceNodeAssembler.containsKey(uuid)) {
return false;
}

Expand All @@ -352,7 +405,9 @@ public boolean loadAll(
}

public boolean deleteAll(String uuid) {
if (!uuid2WriterManager.containsKey(uuid)) {
if (!uuid2WriterManager.containsKey(uuid)
&& !uuid2PieceNodeAssembler.containsKey(uuid)
&& !uuid2CleanupTask.containsKey(uuid)) {
return false;
}
clean(uuid);
Expand All @@ -368,6 +423,7 @@ private void clean(String uuid) {
}
}

uuid2PieceNodeAssembler.remove(uuid);
forceCloseWriterManager(uuid);
}

Expand Down Expand Up @@ -845,6 +901,7 @@ public void run() {
} else {
LOGGER.info(StorageEngineMessages.LOAD_CLEANUP_TASK_STARTS, uuid);
try {
uuid2PieceNodeAssembler.remove(uuid);
forceCloseWriterManager(uuid);
} catch (Exception e) {
LOGGER.warn(StorageEngineMessages.LOAD_CLEANUP_TASK_ERROR, uuid, e);
Expand Down
Loading
Loading