remove refcount (#449)

geaflow/geaflow-core/geaflow-core-common/src/main/java/com/antgroup/geaflow/shuffle/ForwardOutputDesc.java

@@ -30,8 +30,6 @@ public class ForwardOutputDesc<T> implements IOutputDesc, Serializable {
     private final int edgeId;
     // Partition number.
     private final int numPartitions;
-    // Data ref count, for data disposal.
-    private final int refCount;
     // Name of the output edge.
     private final String edgeName;
     // Data exchange mode.
@@ -47,7 +45,6 @@ public ForwardOutputDesc(
         int vertexId,
         int edgeId,
         int numPartitions,
-        int refCount,
         String edgeName,
         DataExchangeMode dataExchangeMode,
         List<Integer> targetTaskIndices,
@@ -56,7 +53,6 @@ public ForwardOutputDesc(
         this.vertexId = vertexId;
         this.edgeId = edgeId;
         this.numPartitions = numPartitions;
-        this.refCount = refCount;
         this.edgeName = edgeName;
         this.dataExchangeMode = dataExchangeMode;
         this.targetTaskIndices = targetTaskIndices;
@@ -76,10 +72,6 @@ public int getNumPartitions() {
         return this.numPartitions;
     }
 
-    public int getRefCount() {
-        return this.refCount;
-    }
-
     public String getEdgeName() {
         return this.edgeName;
     }

geaflow/geaflow-core/geaflow-engine/geaflow-cluster/src/main/java/com/antgroup/geaflow/cluster/collector/PipelineOutputEmitter.java

@@ -112,8 +112,7 @@ public void update(UpdateEmitterRequest request) {
                 .setTaskName(taskArgs.getTaskName())
                 .setChannelNum(forwardOutputDesc.getTargetTaskIndices().size())
                 .setEncoder(encoder)
-                .setDataExchangeMode(forwardOutputDesc.getDataExchangeMode())
-                .setRefCount(forwardOutputDesc.getRefCount());
+                .setDataExchangeMode(forwardOutputDesc.getDataExchangeMode());
             pipeRecordWriter.init(writerContext);
 
             AtomicBoolean flag = new AtomicBoolean(true);

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/api/writer/IWriterContext.java

@@ -87,6 +87,4 @@ public interface IWriterContext extends Serializable {
 
     DataExchangeMode getDataExchangeMode();
 
-    int getRefCount();
-
 }

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/api/writer/PipelineShardWriter.java

@@ -46,8 +46,8 @@ public void init(IWriterContext writerContext) {
     }
 
     @Override
-    protected PipelineSlice newSlice(String taskLogTag, SliceId sliceId, int refCount) {
-        return new PipelineSlice(taskLogTag, sliceId, refCount);
+    protected PipelineSlice newSlice(String taskLogTag, SliceId sliceId) {
+        return new PipelineSlice(taskLogTag, sliceId);
     }
 
     @Override

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/api/writer/ShardWriter.java

@@ -75,7 +75,7 @@ public void init(IWriterContext writerContext) {
         this.maxBufferSize = this.shuffleConfig.getFlushBufferSizeBytes();
 
         this.buffers = this.buildBufferBuilder(this.targetChannels);
-        this.resultSlices = this.buildResultSlices(this.targetChannels, writerContext.getRefCount());
+        this.resultSlices = this.buildResultSlices(this.targetChannels);
         this.recordSerializer = this.getRecordSerializer();
     }
 
@@ -89,20 +89,20 @@ private BufferBuilder[] buildBufferBuilder(int channels) {
         return buffers;
     }
 
-    protected IPipelineSlice[] buildResultSlices(int channels, int refCount) {
+    protected IPipelineSlice[] buildResultSlices(int channels) {
         IPipelineSlice[] slices = new IPipelineSlice[channels];
         WriterId writerId = new WriterId(this.pipelineId, this.edgeId, this.taskIndex);
         SliceManager sliceManager = ShuffleManager.getInstance().getSliceManager();
         for (int i = 0; i < channels; i++) {
             SliceId sliceId = new SliceId(writerId, i);
-            IPipelineSlice slice = this.newSlice(this.taskLogTag, sliceId, refCount);
+            IPipelineSlice slice = this.newSlice(this.taskLogTag, sliceId);
             slices[i] = slice;
             sliceManager.register(sliceId, slice);
         }
         return slices;
     }
 
-    protected abstract IPipelineSlice newSlice(String taskLogTag, SliceId sliceId, int refCount);
+    protected abstract IPipelineSlice newSlice(String taskLogTag, SliceId sliceId);
 
     @SuppressWarnings("unchecked")
     private IRecordSerializer<T> getRecordSerializer() {

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/api/writer/SpillableShardWriter.java

@@ -34,8 +34,8 @@ public SpillableShardWriter(ShuffleAddress shuffleAddress) {
     }
 
     @Override
-    protected IPipelineSlice newSlice(String taskLogTag, SliceId sliceId, int refCount) {
-        return new SpillablePipelineSlice(taskLogTag, sliceId, refCount);
+    protected IPipelineSlice newSlice(String taskLogTag, SliceId sliceId) {
+        return new SpillablePipelineSlice(taskLogTag, sliceId);
     }
 
     @Override

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/api/writer/WriterContext.java

@@ -30,7 +30,6 @@ public class WriterContext implements IWriterContext {
     private DataExchangeMode dataExchangeMode;
     private int targetChannels;
     private ShuffleConfig config;
-    private int refCount;
     private IEncoder<?> encoder;
 
     public WriterContext(long pipelineId, String pipelineName) {
@@ -82,11 +81,6 @@ public WriterContext setEncoder(IEncoder<?> encoder) {
         return this;
     }
 
-    public WriterContext setRefCount(int refCount) {
-        this.refCount = refCount;
-        return this;
-    }
-
     @Override
     public PipelineInfo getPipelineInfo() {
         return pipelineInfo;
@@ -137,11 +131,6 @@ public DataExchangeMode getDataExchangeMode() {
         return this.dataExchangeMode;
     }
 
-    @Override
-    public int getRefCount() {
-        return this.refCount;
-    }
-
     public static WriterContextBuilder newBuilder() {
         return new WriterContextBuilder();
     }

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/network/netty/SliceOutputChannelHandler.java

@@ -254,7 +254,7 @@ public WriteNextMessageIfPossibleListener(PipeBuffer pipeBuffer) {
         @Override
         public void operationComplete(ChannelFuture future) throws Exception {
             try {
-                if (buffer != null && buffer.isDisposable()) {
+                if (buffer != null) {
                     buffer.release();
                 }
                 if (future.isSuccess()) {

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/pipeline/buffer/AbstractBuffer.java

@@ -19,7 +19,6 @@
 public abstract class AbstractBuffer implements OutBuffer {
 
     private final ShuffleMemoryTracker memoryTracker;
-    protected int refCount;
 
     public AbstractBuffer(boolean enableMemoryTrack) {
         this.memoryTracker = enableMemoryTrack
@@ -30,16 +29,6 @@ public AbstractBuffer(ShuffleMemoryTracker memoryTracker) {
         this.memoryTracker = memoryTracker;
     }
 
-    @Override
-    public void setRefCount(int refCount) {
-        this.refCount = refCount;
-    }
-
-    @Override
-    public boolean isDisposable() {
-        return this.refCount <= 0;
-    }
-
     protected void requireMemory(long dataSize) {
         if (this.memoryTracker != null) {
             memoryTracker.requireMemory(dataSize);

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/pipeline/buffer/OutBuffer.java

@@ -50,20 +50,6 @@ public interface OutBuffer {
      */
     void write(OutputStream outputStream) throws IOException;
 
-    /**
-     * Set ref count, the number of consumer which handle this buffer.
-     *
-     * @param refCount ref count.
-     */
-    void setRefCount(int refCount);
-
-    /**
-     * Check if this buffer disposable.
-     *
-     * @return if disposable.
-     */
-    boolean isDisposable();
-
     /**
      * Release this buffer.
      */

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/pipeline/slice/AbstractSlice.java

@@ -27,16 +27,14 @@ public abstract class AbstractSlice implements IPipelineSlice {
 
     protected final SliceId sliceId;
     protected final String taskLogTag;
-    protected int refCount;
     protected int totalBufferCount;
     protected ArrayDeque<PipeBuffer> buffers;
     protected PipelineSliceReader sliceReader;
     protected volatile boolean isReleased;
 
-    public AbstractSlice(String taskLogTag, SliceId sliceId, int refCount) {
+    public AbstractSlice(String taskLogTag, SliceId sliceId) {
         this.sliceId = sliceId;
         this.taskLogTag = taskLogTag;
-        this.refCount = refCount;
         this.totalBufferCount = 0;
         this.buffers = new ArrayDeque<>();
     }
@@ -61,23 +59,17 @@ public PipelineSliceReader createSliceReader(long startBatchId, PipelineSliceLis
                 throw new GeaflowRuntimeException("slice is already created:" + sliceId);
             }
 
-            refCount--;
-            LOGGER.info("creating reader for {} {} with startBatch:{} refCount:{}",
-                taskLogTag, sliceId, startBatchId, refCount);
+            LOGGER.info("creating reader for {} {} with startBatch:{}",
+                taskLogTag, sliceId, startBatchId);
 
-            // multiple repeatable readers can exist at the same time.
-            if (refCount >= 1) {
-                sliceReader = new RepeatableSliceReader(this, startBatchId, listener);
-            } else {
-                sliceReader = new DisposableSliceReader(this, startBatchId, listener);
-            }
+            sliceReader = new DisposableSliceReader(this, startBatchId, listener);
             return sliceReader;
         }
     }
 
     @Override
     public boolean canRelease() {
-        return refCount == 0 && !hasNext();
+        return !hasNext();
     }
 
     @Override

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/pipeline/slice/IPipelineSlice.java

@@ -16,7 +16,6 @@
 
 import com.antgroup.geaflow.shuffle.message.SliceId;
 import com.antgroup.geaflow.shuffle.pipeline.buffer.PipeBuffer;
-import java.util.Iterator;
 
 public interface IPipelineSlice {
 
@@ -60,9 +59,4 @@ public interface IPipelineSlice {
      */
     PipeBuffer next();
 
-    /**
-     * Get slice buffer iterator.
-     */
-    Iterator<PipeBuffer> getBufferIterator();
-
 }

geaflow/geaflow-core/geaflow-engine/geaflow-shuffle/src/main/java/com/antgroup/geaflow/shuffle/pipeline/slice/PipelineSlice.java

@@ -17,7 +17,6 @@
 import com.antgroup.geaflow.shuffle.message.SliceId;
 import com.antgroup.geaflow.shuffle.pipeline.buffer.PipeBuffer;
 import com.google.common.base.Preconditions;
-import java.util.Iterator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -28,11 +27,7 @@ public class PipelineSlice extends AbstractSlice {
     private boolean flushRequested;
 
     public PipelineSlice(String taskLogTag, SliceId sliceId) {
-        super(taskLogTag, sliceId, 1);
-    }
-
-    public PipelineSlice(String taskLogTag, SliceId sliceId, int refCount) {
-        super(taskLogTag, sliceId, refCount);
+        super(taskLogTag, sliceId);
     }
 
     // ------------------------------------------------------------------------
@@ -87,12 +82,6 @@ private void notifyDataAvailable(long batchId) {
         }
     }
 
-    public void setRefCount(int refCount) {
-        synchronized (buffers) {
-            this.refCount = refCount;
-        }
-    }
-
     // ------------------------------------------------------------------------
     // Consume
     // ------------------------------------------------------------------------
@@ -133,11 +122,6 @@ private void updateFlushRequested(boolean flushRequested) {
         this.flushRequested = flushRequested;
     }
 
-    @Override
-    public Iterator<PipeBuffer> getBufferIterator() {
-        return buffers.iterator();
-    }
-
     @Override
     public void release() {
         final PipelineSliceReader reader;