[Improve][API & Zeta] Using connector custom serializer encode/decode states

seatunnel-api/src/main/java/org/apache/seatunnel/api/serialization/DefaultSerializer.java

@@ -35,6 +35,9 @@ public byte[] serialize(T obj) throws IOException {
 
     @Override
     public T deserialize(byte[] serialized) throws IOException {
+        if (serialized == null) {
+            return null;
+        }
         return SerializationUtils.deserialize(serialized);
     }
 }

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalPlanGenerator.java

@@ -552,7 +552,10 @@ private List<PhysicalVertex> getSourceTask(
                                                                                 .getJobId(),
                                                                         taskLocation,
                                                                         finalParallelismIndex,
-                                                                        f);
+                                                                        (PhysicalExecutionFlow<
+                                                                                        SourceAction,
+                                                                                        SourceConfig>)
+                                                                                f);
                                                             } else {
                                                                 return new TransformSeaTunnelTask(
                                                                         jobImmutableInformation

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/SinkAggregatedCommitterTask.java

@@ -34,6 +34,7 @@
 import org.apache.commons.collections4.CollectionUtils;
 
 import com.hazelcast.cluster.Address;
+import lombok.Getter;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 
@@ -45,6 +46,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
@@ -75,6 +77,8 @@ public class SinkAggregatedCommitterTask<CommandInfoT, AggregatedCommitInfoT>
     private final SinkAggregatedCommitter<CommandInfoT, AggregatedCommitInfoT> aggregatedCommitter;
 
     private transient Serializer<AggregatedCommitInfoT> aggregatedCommitInfoSerializer;
+    @Getter private transient Serializer<CommandInfoT> commitInfoSerializer;
+
     private Map<Long, Address> writerAddressMap;
 
     private ConcurrentMap<Long, List<CommandInfoT>> commitInfoCache;
@@ -107,6 +111,7 @@ public void init() throws Exception {
         this.writerAddressMap = new ConcurrentHashMap<>();
         this.checkpointCommitInfoMap = new ConcurrentHashMap<>();
         this.completableFuture = new CompletableFuture<>();
+        this.commitInfoSerializer = sink.getSink().getCommitInfoSerializer().get();
         this.aggregatedCommitInfoSerializer =
                 sink.getSink().getAggregatedCommitInfoSerializer().get();
         log.debug(
@@ -250,6 +255,7 @@ public void restoreState(List<ActionSubtaskState> actionStateList) throws Except
                 actionStateList.stream()
                         .map(ActionSubtaskState::getState)
                         .flatMap(Collection::stream)
+                        .filter(Objects::nonNull)
                         .map(
                                 bytes ->
                                         sneaky(

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/SourceSeaTunnelTask.java

@@ -18,17 +18,19 @@
 package org.apache.seatunnel.engine.server.task;
 
 import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.serialization.Serializer;
 import org.apache.seatunnel.api.source.SourceSplit;
 import org.apache.seatunnel.engine.core.dag.actions.SourceAction;
 import org.apache.seatunnel.engine.server.dag.physical.config.SourceConfig;
-import org.apache.seatunnel.engine.server.dag.physical.flow.Flow;
+import org.apache.seatunnel.engine.server.dag.physical.flow.PhysicalExecutionFlow;
 import org.apache.seatunnel.engine.server.execution.ProgressState;
 import org.apache.seatunnel.engine.server.execution.TaskLocation;
 import org.apache.seatunnel.engine.server.task.flow.SourceFlowLifeCycle;
 import org.apache.seatunnel.engine.server.task.record.Barrier;
 
 import com.hazelcast.logging.ILogger;
 import com.hazelcast.logging.Logger;
+import lombok.Getter;
 import lombok.NonNull;
 
 import java.util.List;
@@ -41,15 +43,24 @@ public class SourceSeaTunnelTask<T, SplitT extends SourceSplit> extends SeaTunne
     private transient SeaTunnelSourceCollector<T> collector;
 
     private transient Object checkpointLock;
+    @Getter private transient Serializer<SplitT> splitSerializer;
+    private final PhysicalExecutionFlow<SourceAction, SourceConfig> sourceFlow;
 
-    public SourceSeaTunnelTask(long jobID, TaskLocation taskID, int indexID, Flow executionFlow) {
+    public SourceSeaTunnelTask(
+            long jobID,
+            TaskLocation taskID,
+            int indexID,
+            PhysicalExecutionFlow<SourceAction, SourceConfig> executionFlow) {
         super(jobID, taskID, indexID, executionFlow);
+        this.sourceFlow = executionFlow;
     }
 
     @Override
     public void init() throws Exception {
         super.init();
         this.checkpointLock = new Object();
+        this.splitSerializer = sourceFlow.getAction().getSource().getSplitSerializer();
+
         LOGGER.info("starting seatunnel source task, index " + indexID);
         if (!(startFlowLifeCycle instanceof SourceFlowLifeCycle)) {
             throw new TaskRuntimeException(

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTask.java

@@ -37,6 +37,7 @@
 import com.hazelcast.cluster.Address;
 import com.hazelcast.spi.impl.operationservice.Operation;
 import com.hazelcast.spi.impl.operationservice.impl.InvocationFuture;
+import lombok.Getter;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 
@@ -77,6 +78,7 @@ public class SourceSplitEnumeratorTask<SplitT extends SourceSplit> extends Coord
     private SeaTunnelSplitEnumeratorContext<SplitT> enumeratorContext;
 
     private Serializer<Serializable> enumeratorStateSerializer;
+    @Getter private Serializer<SplitT> splitSerializer;
 
     private int maxReaderSize;
     private Set<Long> unfinishedReaders;
@@ -102,6 +104,7 @@ public void init() throws Exception {
                 new SeaTunnelSplitEnumeratorContext<>(
                         this.source.getParallelism(), this, getMetricsContext());
         enumeratorStateSerializer = this.source.getSource().getEnumeratorStateSerializer();
+        splitSerializer = this.source.getSource().getSplitSerializer();
         taskMemberMapping = new ConcurrentHashMap<>();
         taskIDToTaskLocationMapping = new ConcurrentHashMap<>();
         taskIndexToTaskLocationMapping = new ConcurrentHashMap<>();

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/context/SeaTunnelSplitEnumeratorContext.java

@@ -21,7 +21,6 @@
 import org.apache.seatunnel.api.source.SourceEvent;
 import org.apache.seatunnel.api.source.SourceSplit;
 import org.apache.seatunnel.api.source.SourceSplitEnumerator;
-import org.apache.seatunnel.common.utils.SerializationUtils;
 import org.apache.seatunnel.engine.server.task.SourceSplitEnumeratorTask;
 import org.apache.seatunnel.engine.server.task.operation.source.AssignSplitOperation;
 
@@ -31,6 +30,9 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.apache.seatunnel.engine.common.utils.ExceptionUtil.sneaky;
 
 @Slf4j
 public class SeaTunnelSplitEnumeratorContext<SplitT extends SourceSplit>
@@ -67,22 +69,26 @@ public void assignSplit(int subtaskIndex, List<SplitT> splits) {
             log.warn("No reader is obtained, skip this assign!");
             return;
         }
+
+        List<byte[]> splitBytes =
+                splits.stream()
+                        .map(split -> sneaky(() -> task.getSplitSerializer().serialize(split)))
+                        .collect(Collectors.toList());
         task.getExecutionContext()
                 .sendToMember(
                         new AssignSplitOperation<>(
-                                task.getTaskMemberLocationByIndex(subtaskIndex),
-                                SerializationUtils.serialize(splits.toArray())),
+                                task.getTaskMemberLocationByIndex(subtaskIndex), splitBytes),
                         task.getTaskMemberAddressByIndex(subtaskIndex))
                 .join();
     }
 
     @Override
     public void signalNoMoreSplits(int subtaskIndex) {
+        List<byte[]> emptySplits = Collections.emptyList();
         task.getExecutionContext()
                 .sendToMember(
                         new AssignSplitOperation<>(
-                                task.getTaskMemberLocationByIndex(subtaskIndex),
-                                SerializationUtils.serialize(Collections.emptyList().toArray())),
+                                task.getTaskMemberLocationByIndex(subtaskIndex), emptySplits),
                         task.getTaskMemberAddressByIndex(subtaskIndex))
                 .join();
     }

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SinkFlowLifeCycle.java

@@ -25,7 +25,6 @@
 import org.apache.seatunnel.api.sink.SinkWriter;
 import org.apache.seatunnel.api.table.event.SchemaChangeEvent;
 import org.apache.seatunnel.api.table.type.Record;
-import org.apache.seatunnel.common.utils.SerializationUtils;
 import org.apache.seatunnel.engine.core.checkpoint.InternalCheckpointListener;
 import org.apache.seatunnel.engine.core.dag.actions.SinkAction;
 import org.apache.seatunnel.engine.server.checkpoint.ActionStateKey;
@@ -48,6 +47,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -66,6 +66,7 @@ public class SinkFlowLifeCycle<T, CommitInfoT extends Serializable, AggregatedCo
     private final SinkAction<T, StateT, CommitInfoT, AggregatedCommitInfoT> sinkAction;
     private SinkWriter<T, CommitInfoT, StateT> writer;
 
+    private transient Optional<Serializer<CommitInfoT>> commitInfoSerializer;
     private transient Optional<Serializer<StateT>> writerStateSerializer;
 
     private final int indexID;
@@ -110,6 +111,7 @@ public SinkFlowLifeCycle(
 
     @Override
     public void init() throws Exception {
+        this.commitInfoSerializer = sinkAction.getSink().getCommitInfoSerializer();
         this.writerStateSerializer = sinkAction.getSink().getWriterStateSerializer();
         this.committer = sinkAction.getSink().createCommitter();
         this.lastCommitInfo = Optional.empty();
@@ -184,10 +186,14 @@ public void received(Record<?> record) {
                         runningTask
                                 .getExecutionContext()
                                 .sendToMember(
-                                        new SinkPrepareCommitOperation(
+                                        new SinkPrepareCommitOperation<CommitInfoT>(
                                                 barrier,
                                                 committerTaskLocation,
-                                                SerializationUtils.serialize(commitInfoT)),
+                                                commitInfoSerializer.isPresent()
+                                                        ? commitInfoSerializer
+                                                                .get()
+                                                                .serialize(commitInfoT)
+                                                        : null),
                                         committerTaskAddress)
                                 .join();
                     }
@@ -247,9 +253,9 @@ public void restoreState(List<ActionSubtaskState> actionStateList) throws Except
         if (writerStateSerializer.isPresent()) {
             states =
                     actionStateList.stream()
-                            .filter(state -> writerStateSerializer.isPresent())
                             .map(ActionSubtaskState::getState)
                             .flatMap(Collection::stream)
+                            .filter(Objects::nonNull)
                             .map(
                                     bytes ->
                                             sneaky(

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceFlowLifeCycle.java

@@ -23,7 +23,6 @@
 import org.apache.seatunnel.api.source.SourceReader;
 import org.apache.seatunnel.api.source.SourceSplit;
 import org.apache.seatunnel.api.table.type.Record;
-import org.apache.seatunnel.common.utils.SerializationUtils;
 import org.apache.seatunnel.engine.core.checkpoint.CheckpointType;
 import org.apache.seatunnel.engine.core.checkpoint.InternalCheckpointListener;
 import org.apache.seatunnel.engine.core.dag.actions.SourceAction;
@@ -59,7 +58,6 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
-import static org.apache.seatunnel.engine.common.utils.ExceptionUtil.sneaky;
 import static org.apache.seatunnel.engine.server.task.AbstractTask.serializeStates;
 
 @Slf4j
@@ -338,21 +336,17 @@ public void restoreState(List<ActionSubtaskState> actionStateList) throws Except
         if (actionStateList.isEmpty()) {
             return;
         }
-        List<SplitT> splits =
+        List<byte[]> splits =
                 actionStateList.stream()
                         .map(ActionSubtaskState::getState)
                         .flatMap(Collection::stream)
                         .filter(Objects::nonNull)
-                        .map(bytes -> sneaky(() -> splitSerializer.deserialize(bytes)))
                         .collect(Collectors.toList());
         try {
             runningTask
                     .getExecutionContext()
                     .sendToMember(
-                            new RestoredSplitOperation(
-                                    enumeratorTaskLocation,
-                                    SerializationUtils.serialize(splits.toArray()),
-                                    indexID),
+                            new RestoredSplitOperation(enumeratorTaskLocation, splits, indexID),
                             enumeratorTaskAddress)
                     .get();
         } catch (InterruptedException | ExecutionException e) {

seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/operation/sink/SinkPrepareCommitOperation.java

@@ -17,7 +17,6 @@
 
 package org.apache.seatunnel.engine.server.task.operation.sink;
 
-import org.apache.seatunnel.common.utils.SerializationUtils;
 import org.apache.seatunnel.engine.server.SeaTunnelServer;
 import org.apache.seatunnel.engine.server.TaskExecutionService;
 import org.apache.seatunnel.engine.server.execution.TaskLocation;
@@ -33,7 +32,7 @@
 import java.io.IOException;
 
 @NoArgsConstructor
-public class SinkPrepareCommitOperation extends BarrierFlowOperation {
+public class SinkPrepareCommitOperation<CommitInfoT> extends BarrierFlowOperation {
     private byte[] commitInfos;
 
     public SinkPrepareCommitOperation(
@@ -73,15 +72,24 @@ public int getClassId() {
     public void run() throws Exception {
         TaskExecutionService taskExecutionService =
                 ((SeaTunnelServer) getService()).getTaskExecutionService();
-        SinkAggregatedCommitterTask<?, ?> committerTask =
+        SinkAggregatedCommitterTask<CommitInfoT, ?> committerTask =
                 taskExecutionService.getTask(taskLocation);
-        ClassLoader classLoader =
+        ClassLoader taskClassLoader =
                 taskExecutionService
                         .getExecutionContext(taskLocation.getTaskGroupLocation())
                         .getClassLoader();
+        ClassLoader mainClassLoader = Thread.currentThread().getContextClassLoader();
+
         if (commitInfos != null) {
-            committerTask.receivedWriterCommitInfo(
-                    barrier.getId(), SerializationUtils.deserialize(commitInfos, classLoader));
+            CommitInfoT deserializeCommitInfo = null;
+            try {
+                Thread.currentThread().setContextClassLoader(taskClassLoader);
+                deserializeCommitInfo =
+                        committerTask.getCommitInfoSerializer().deserialize(commitInfos);
+            } finally {
+                Thread.currentThread().setContextClassLoader(mainClassLoader);
+            }
+            committerTask.receivedWriterCommitInfo(barrier.getId(), deserializeCommitInfo);
         }
         committerTask.triggerBarrier(barrier);
     }