[#436] feat(client,server): Introduce multi-part LocalStorageManager #2253

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.

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

Merged

maobaolong merged 8 commits into apache:master from maobaolong:raidStorage.gh

Nov 26, 2024

common/src/main/java/io/netty/util/CompositeFileRegion.java

-Original file line number
+Diff line change
@@ -0,0 +1,142 @@
+    /*
+     * 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 io.netty.util;
+    import java.io.IOException;
+    import java.nio.channels.WritableByteChannel;
+    import io.netty.channel.FileRegion;
+    import org.apache.uniffle.common.netty.protocol.AbstractFileRegion;
+    public class CompositeFileRegion extends AbstractFileRegion {
+      private final FileRegion[] regions;
+      private long totalSize = 0;
+      private long bytesTransferred = 0;
+      public CompositeFileRegion(FileRegion... regions) {
+        this.regions = regions;
+        for (FileRegion region : regions) {
+          totalSize += region.count();
+        }
+      }
+      @Override
+      public long position() {
+        return bytesTransferred;
+      }
+      @Override
+      public long count() {
+        return totalSize;
+      }
+      @Override
+      public long transferTo(WritableByteChannel target, long position) throws IOException {
+        long totalBytesTransferred = 0;
+        for (FileRegion region : regions) {
+          if (position >= region.count()) {
+            position -= region.count();
+          } else {
+            long currentBytesTransferred = region.transferTo(target, position);
+            totalBytesTransferred += currentBytesTransferred;
+            bytesTransferred += currentBytesTransferred;
+            if (currentBytesTransferred < region.count() - position) {
+              break;
+            }
+            position = 0;
+          }
+        }
+        return totalBytesTransferred;
+      }
+      @Override
+      public long transferred() {
+        return bytesTransferred;
+      }
+      @Override
+      public AbstractFileRegion retain() {
+        super.retain();
+        for (FileRegion region : regions) {
+          region.retain();
+        }
+        return this;
+      }
+      @Override
+      public AbstractFileRegion retain(int increment) {
+        super.retain(increment);
+        for (FileRegion region : regions) {
+          region.retain(increment);
+        }
+        return this;
+      }
+      @Override
+      public boolean release() {
+        boolean released = super.release();
+        for (FileRegion region : regions) {
+          if (!region.release()) {
+            released = false;
+          }
+        }
+        return released;
+      }
+      @Override
+      public boolean release(int decrement) {
+        boolean released = super.release(decrement);
+        for (FileRegion region : regions) {
+          if (!region.release(decrement)) {
+            released = false;
+          }
+        }
+        return released;
+      }
+      @Override
+      protected void deallocate() {
+        for (FileRegion region : regions) {
+          if (region instanceof AbstractReferenceCounted) {
+            ((AbstractReferenceCounted) region).deallocate();
+          }
+        }
+      }
+      @Override
+      public AbstractFileRegion touch() {
+        super.touch();
+        for (FileRegion region : regions) {
+          region.touch();
+        }
+        return this;
+      }
+      @Override
+      public AbstractFileRegion touch(Object hint) {
+        super.touch(hint);
+        for (FileRegion region : regions) {
+          region.touch(hint);
+        }
+        return this;
+      }
+    }

common/src/main/java/org/apache/uniffle/common/ShuffleDataSegment.java

-Original file line number
+Diff line change
@@ Expand Up / @@ -27,11 +27,15 @@ @@
     public class ShuffleDataSegment {
       private final long offset;
       private final int length;
+      private final int storageId;
       private final List<BufferSegment> bufferSegments;
-      public ShuffleDataSegment(long offset, int length, List<BufferSegment> bufferSegments) {
+      public ShuffleDataSegment(
+          long offset, int length, int storageId, List<BufferSegment> bufferSegments) {
         this.offset = offset;
         this.length = length;
+        this.storageId = storageId;
         this.bufferSegments = bufferSegments;
       }
@@ Expand All / @@ -46,4 +50,8 @@ public int getLength() { @@
       public List<BufferSegment> getBufferSegments() {
         return bufferSegments;
       }
+      public int getStorageId() {
+        return storageId;
+      }
     }

common/src/main/java/org/apache/uniffle/common/ShuffleIndexResult.java

-Original file line number
+Diff line change
@@ Expand Up / @@ -30,8 +30,10 @@ @@
     public class ShuffleIndexResult {
       private static final Logger LOG = LoggerFactory.getLogger(ShuffleIndexResult.class);
+      private static final int[] DEFAULT_STORAGE_IDS = new int[] {0};
       private final ManagedBuffer buffer;
+      private final int[] storageIds;
       private long dataFileLen;
       private String dataFileName;
@@ Expand All / @@ -44,15 +46,28 @@ public ShuffleIndexResult(byte[] data, long dataFileLen) { @@
       }
       public ShuffleIndexResult(ByteBuffer data, long dataFileLen) {
-        this.buffer =
-            new NettyManagedBuffer(data != null ? Unpooled.wrappedBuffer(data) : Unpooled.EMPTY_BUFFER);
-        this.dataFileLen = dataFileLen;
+        this(
+            new NettyManagedBuffer(data != null ? Unpooled.wrappedBuffer(data) : Unpooled.EMPTY_BUFFER),
+            dataFileLen,
+            null,
+            DEFAULT_STORAGE_IDS);
       }
       public ShuffleIndexResult(ManagedBuffer buffer, long dataFileLen, String dataFileName) {
+        this(buffer, dataFileLen, dataFileName, DEFAULT_STORAGE_IDS);
+      }
+      public ShuffleIndexResult(
+          ManagedBuffer buffer, long dataFileLen, String dataFileName, int storageId) {
+        this(buffer, dataFileLen, dataFileName, new int[] {storageId});
+      }
+      public ShuffleIndexResult(
+          ManagedBuffer buffer, long dataFileLen, String dataFileName, int[] storageIds) {
         this.buffer = buffer;
         this.dataFileLen = dataFileLen;
         this.dataFileName = dataFileName;
+        this.storageIds = storageIds;
       }
       public byte[] getData() {
@@ Expand Down Expand Up / @@ -99,4 +114,8 @@ public ManagedBuffer getManagedBuffer() { @@
       public String getDataFileName() {
         return dataFileName;
       }
+      public int[] getStorageIds() {
+        return storageIds;
+      }
     }

common/src/main/java/org/apache/uniffle/common/netty/MessageEncoder.java

-Original file line number
+Diff line change
@@ Expand Up @@
         header.writeInt(bodyLength);
         in.encode(header);
         if (header.writableBytes() != 0) {
-          throw new RssException("header's writable bytes should be 0");
+          throw new RssException(
+              "header's writable bytes should be 0, but it is " + header.writableBytes());
         }
         if (body != null) {
@@ Expand Down @@

...n/src/main/java/org/apache/uniffle/common/netty/buffer/MultiFileSegmentManagedBuffer.java

-Original file line number
+Diff line change
@@ -0,0 +1,83 @@
+    /*
+     * 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.uniffle.common.netty.buffer;
+    import java.nio.ByteBuffer;
+    import java.util.ArrayList;
+    import java.util.List;
+    import io.netty.buffer.ByteBuf;
+    import io.netty.buffer.Unpooled;
+    import io.netty.channel.FileRegion;
+    import io.netty.util.CompositeFileRegion;
+    import org.slf4j.Logger;
+    import org.slf4j.LoggerFactory;
+    /** A wrapper of multiple {@link FileSegmentManagedBuffer}, used for combine shuffle index files. */
+    public class MultiFileSegmentManagedBuffer extends ManagedBuffer {
+      private static final Logger LOG = LoggerFactory.getLogger(MultiFileSegmentManagedBuffer.class);
+      private final List<ManagedBuffer> managedBuffers;
+      public MultiFileSegmentManagedBuffer(List<ManagedBuffer> managedBuffers) {
+        this.managedBuffers = managedBuffers;
+      }
+      @Override
+      public int size() {
+        return managedBuffers.stream().mapToInt(ManagedBuffer::size).sum();
+      }
+      @Override
+      public ByteBuf byteBuf() {
+        return Unpooled.wrappedBuffer(this.nioByteBuffer());
+      }
+      @Override
+      public ByteBuffer nioByteBuffer() {
+        ByteBuffer merged = ByteBuffer.allocate(size());
+        for (ManagedBuffer managedBuffer : managedBuffers) {
+          ByteBuffer buffer = managedBuffer.nioByteBuffer();
+          merged.put(buffer.slice());
+        }
+        merged.flip();
+        return merged;
+      }
+      @Override
+      public ManagedBuffer retain() {
+        return this;
+      }
+      @Override
+      public ManagedBuffer release() {
+        return this;
+      }
+      @Override
+      public Object convertToNetty() {
+        List<FileRegion> fileRegions = new ArrayList<>(managedBuffers.size());
+        for (ManagedBuffer managedBuffer : managedBuffers) {
+          Object object = managedBuffer.convertToNetty();
+          if (object instanceof FileRegion) {
+            fileRegions.add((FileRegion) object);
+          }
+        }
+        return new CompositeFileRegion(fileRegions.toArray(new FileRegion[0]));
+      }
+    }

...on/src/main/java/org/apache/uniffle/common/netty/protocol/GetLocalShuffleDataRequest.java

-Original file line number
+Diff line change
@@ Expand Up @@
       private long offset;
       private int length;
       private long timestamp;
+      private int storageId;
       public GetLocalShuffleDataRequest(
           long requestId,
@@ Expand All / @@ -41,6 +42,30 @@ public GetLocalShuffleDataRequest( @@
           long offset,
           int length,
           long timestamp) {
+        this(
+            requestId,
+            appId,
+            shuffleId,
+            partitionId,
+            partitionNumPerRange,
+            partitionNum,
+            offset,
+            length,
+            -1,
+            timestamp);
+      }
+      protected GetLocalShuffleDataRequest(
+          long requestId,
+          String appId,
+          int shuffleId,
+          int partitionId,
+          int partitionNumPerRange,
+          int partitionNum,
+          long offset,
+          int length,
+          int storageId,
+          long timestamp) {
         super(requestId);
         this.appId = appId;
         this.shuffleId = shuffleId;
@@ Expand All / @@ -49,6 +74,7 @@ public GetLocalShuffleDataRequest( @@
         this.partitionNum = partitionNum;
         this.offset = offset;
         this.length = length;
+        this.storageId = storageId;
         this.timestamp = timestamp;
       }
@@ Expand Down Expand Up / @@ -132,6 +158,10 @@ public long getTimestamp() { @@
         return timestamp;
       }
+      public int getStorageId() {
+        return storageId;
+      }
       @Override
       public String getOperationType() {
         return "getLocalShuffleData";
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[#436] feat(client,server): Introduce multi-part LocalStorageManager #2253

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!