[fix](routineload) Preserve paused job metadata during replay

fe/fe-core/src/main/java/org/apache/doris/load/routineload/RoutineLoadJob.java

@@ -1514,13 +1514,36 @@ protected void unprotectUpdateState(JobState jobState, ErrorReason reason, boole
 
         if (!isReplay && jobState != JobState.RUNNING) {
             if (jobState == JobState.PAUSED) {
-                Env.getCurrentEnv().getEditLog().logOpRoutineLoadJob(new RoutineLoadOperation(id, jobState, reason));
+                Env.getCurrentEnv().getEditLog().logOpRoutineLoadJob(createPauseOperation(reason));
             } else {
                 Env.getCurrentEnv().getEditLog().logOpRoutineLoadJob(new RoutineLoadOperation(id, jobState));
             }
         }
     }
 
+    protected RoutineLoadOperation createPauseOperation(ErrorReason reason) {
+        return new RoutineLoadOperation(id, JobState.PAUSED, reason);
+    }
+
+    public void replayRestoreOperatorMetadata(RoutineLoadOperation operation) {
+        if (operation.getProgress() == null) {
+            // Pause operations written before operator metadata was persisted do not contain a snapshot.
+            return;
+        }
+        Preconditions.checkNotNull(operation.getStatistic());
+        writeLock();
+        try {
+            progress = operation.getProgress();
+            jobStatistic = operation.getStatistic();
+            replayRestoreDataSourceOperatorMetadata(operation);
+        } finally {
+            writeUnlock();
+        }
+    }
+
+    protected void replayRestoreDataSourceOperatorMetadata(RoutineLoadOperation operation) {
+    }
+
     private void executeRunning() {
         state = JobState.RUNNING;
     }

fe/fe-core/src/main/java/org/apache/doris/load/routineload/RoutineLoadManager.java

@@ -908,6 +908,7 @@ public void replayChangeRoutineLoadJob(RoutineLoadOperation operation) {
         RoutineLoadJob job = getJob(operation.getId());
         try {
             job.updateState(operation.getJobState(), operation.getErrorReason(), true /* is replay */);
+            job.replayRestoreOperatorMetadata(operation);
         } catch (UserException e) {
             LOG.error("should not happened", e);
         } catch (NullPointerException npe) {

fe/fe-core/src/main/java/org/apache/doris/load/routineload/kafka/KafkaRoutineLoadJob.java

@@ -56,6 +56,7 @@
 import org.apache.doris.nereids.trees.plans.commands.AlterRoutineLoadCommand;
 import org.apache.doris.nereids.trees.plans.commands.info.CreateRoutineLoadInfo;
 import org.apache.doris.persist.AlterRoutineLoadJobOperationLog;
+import org.apache.doris.persist.RoutineLoadOperation;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.rpc.RpcException;
 import org.apache.doris.service.FrontendOptions;
@@ -78,6 +79,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -114,6 +116,7 @@ public class KafkaRoutineLoadJob extends RoutineLoadJob {
     @SerializedName("cskp")
     private List<Integer> customKafkaPartitions = Lists.newArrayList();
     // current kafka partitions is the actual partition which will be fetched
+    @SerializedName("ckp")
     private List<Integer> currentKafkaPartitions = Lists.newArrayList();
     // optional, user want to set default offset when new partition add or offset not set.
     // kafkaDefaultOffSet has two formats, one is the time format, eg: "2021-10-10 11:00:00",
@@ -127,6 +130,7 @@ public class KafkaRoutineLoadJob extends RoutineLoadJob {
 
     // The latest offset of each partition fetched from kafka server.
     // Will be updated periodically by calling hasMoreDataToConsume()
+    @SerializedName("cplo")
     private Map<Integer, Long> cachedPartitionWithLatestOffsets = Maps.newConcurrentMap();
 
     // The kafka partition fetch from kafka server.
@@ -157,6 +161,14 @@ public KafkaRoutineLoadJob(Long id, String name,
         setMultiTable(isMultiTable);
     }
 
+    @Override
+    public void gsonPostProcess() throws IOException {
+        super.gsonPostProcess();
+        Map<Integer, Long> replayedLatestOffsets = Maps.newConcurrentMap();
+        replayedLatestOffsets.putAll(Preconditions.checkNotNull(cachedPartitionWithLatestOffsets));
+        cachedPartitionWithLatestOffsets = replayedLatestOffsets;
+    }
+
     public String getTopic() {
         return topic;
     }
@@ -543,6 +555,21 @@ protected boolean needAutoResume() {
         }
     }
 
+    @Override
+    protected RoutineLoadOperation createPauseOperation(ErrorReason reason) {
+        return new RoutineLoadOperation(id, JobState.PAUSED, reason, progress, jobStatistic,
+                currentKafkaPartitions, cachedPartitionWithLatestOffsets);
+    }
+
+    @Override
+    protected void replayRestoreDataSourceOperatorMetadata(RoutineLoadOperation operation) {
+        currentKafkaPartitions = Lists.newArrayList(
+                Preconditions.checkNotNull(operation.getCurrentKafkaPartitions()));
+        cachedPartitionWithLatestOffsets = Maps.newConcurrentMap();
+        cachedPartitionWithLatestOffsets.putAll(
+                Preconditions.checkNotNull(operation.getCachedPartitionWithLatestOffsets()));
+    }
+
     @Override
     public String getStatistic() {
         Map<String, Object> summary = this.jobStatistic.summary();

fe/fe-core/src/main/java/org/apache/doris/persist/RoutineLoadOperation.java

@@ -21,13 +21,19 @@
 import org.apache.doris.common.io.Writable;
 import org.apache.doris.load.routineload.ErrorReason;
 import org.apache.doris.load.routineload.RoutineLoadJob.JobState;
+import org.apache.doris.load.routineload.RoutineLoadProgress;
+import org.apache.doris.load.routineload.RoutineLoadStatistic;
 import org.apache.doris.persist.gson.GsonUtils;
 
 import com.google.gson.annotations.SerializedName;
 
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 public class RoutineLoadOperation implements Writable {
     @SerializedName("id")
@@ -36,6 +42,14 @@ public class RoutineLoadOperation implements Writable {
     private JobState jobState;
     @SerializedName("rs")
     private ErrorReason reason;
+    @SerializedName("opg")
+    private RoutineLoadProgress progress;
+    @SerializedName("ost")
+    private RoutineLoadStatistic statistic;
+    @SerializedName("ckp")
+    private List<Integer> currentKafkaPartitions;
+    @SerializedName("cplo")
+    private Map<Integer, Long> cachedPartitionWithLatestOffsets;
 
     private RoutineLoadOperation() {
     }
@@ -51,6 +65,16 @@ public RoutineLoadOperation(long id, JobState jobState, ErrorReason reason) {
         this.reason = reason;
     }
 
+    public RoutineLoadOperation(long id, JobState jobState, ErrorReason reason,
+            RoutineLoadProgress progress, RoutineLoadStatistic statistic,
+            List<Integer> currentKafkaPartitions, Map<Integer, Long> cachedPartitionWithLatestOffsets) {
+        this(id, jobState, reason);
+        this.progress = progress;
+        this.statistic = statistic;
+        this.currentKafkaPartitions = new ArrayList<>(currentKafkaPartitions);
+        this.cachedPartitionWithLatestOffsets = new HashMap<>(cachedPartitionWithLatestOffsets);
+    }
+
     public long getId() {
         return id;
     }
@@ -63,6 +87,22 @@ public ErrorReason getErrorReason() {
         return reason;
     }
 
+    public RoutineLoadProgress getProgress() {
+        return progress;
+    }
+
+    public RoutineLoadStatistic getStatistic() {
+        return statistic;
+    }
+
+    public List<Integer> getCurrentKafkaPartitions() {
+        return currentKafkaPartitions;
+    }
+
+    public Map<Integer, Long> getCachedPartitionWithLatestOffsets() {
+        return cachedPartitionWithLatestOffsets;
+    }
+
     public static RoutineLoadOperation read(DataInput in) throws IOException {
         return GsonUtils.GSON.fromJson(Text.readString(in), RoutineLoadOperation.class);
     }

fe/fe-core/src/test/java/org/apache/doris/load/routineload/KafkaRoutineLoadJobTest.java

@@ -26,6 +26,7 @@
 import org.apache.doris.catalog.Table;
 import org.apache.doris.catalog.info.PartitionNamesInfo;
 import org.apache.doris.common.Config;
+import org.apache.doris.common.InternalErrorCode;
 import org.apache.doris.common.MetaNotFoundException;
 import org.apache.doris.common.Pair;
 import org.apache.doris.common.UserException;
@@ -44,6 +45,9 @@
 import org.apache.doris.nereids.trees.plans.commands.info.LabelNameInfo;
 import org.apache.doris.nereids.trees.plans.commands.load.LoadProperty;
 import org.apache.doris.nereids.trees.plans.commands.load.LoadSeparator;
+import org.apache.doris.persist.EditLog;
+import org.apache.doris.persist.RoutineLoadOperation;
+import org.apache.doris.persist.gson.GsonUtils;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.thrift.TResourceInfo;
 import org.apache.doris.thrift.TRoutineLoadTask;
@@ -58,9 +62,14 @@
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -248,6 +257,110 @@ public void testUpdateLagRebuildsConvertedPropertiesAfterReplay() throws UserExc
         }
     }
 
+    @Test
+    public void testReplayPausedJobPreservesOperatorMetadata() throws Exception {
+        long jobId = 1L;
+        KafkaRoutineLoadJob sourceJob = new KafkaRoutineLoadJob(jobId, "kafka_routine_load_job", 1L,
+                1L, "127.0.0.1:9020", "topic1", UserIdentity.ADMIN);
+
+        Map<Integer, Long> partitionOffsets = Maps.newHashMap();
+        partitionOffsets.put(1, 11L);
+        partitionOffsets.put(2, 21L);
+        Deencapsulation.setField(sourceJob, "progress", new KafkaProgress(partitionOffsets));
+
+        RoutineLoadStatistic statistic = sourceJob.getRoutineLoadStatistic();
+        statistic.totalRows = 100L;
+        statistic.errorRows = 2L;
+        statistic.receivedBytes = 1024L;
+
+        List<Integer> currentKafkaPartitions = Lists.newArrayList(1, 2);
+        Map<Integer, Long> latestOffsets = Maps.newHashMap();
+        latestOffsets.put(1, 15L);
+        latestOffsets.put(2, 30L);
+        Deencapsulation.setField(sourceJob, "currentKafkaPartitions", currentKafkaPartitions);
+        Deencapsulation.setField(sourceJob, "cachedPartitionWithLatestOffsets", latestOffsets);
+
+        ErrorReason pauseReason = new ErrorReason(InternalErrorCode.MANUAL_PAUSE_ERR, "pause for replay");
+        RoutineLoadOperation pauseOperation;
+        Env env = Mockito.mock(Env.class);
+        EditLog editLog = Mockito.mock(EditLog.class);
+        try (MockedStatic<Env> envStatic = Mockito.mockStatic(Env.class)) {
+            envStatic.when(Env::getCurrentEnv).thenReturn(env);
+            Mockito.when(env.getEditLog()).thenReturn(editLog);
+            sourceJob.updateState(RoutineLoadJob.JobState.PAUSED, pauseReason, false);
+            ArgumentCaptor<RoutineLoadOperation> operationCaptor =
+                    ArgumentCaptor.forClass(RoutineLoadOperation.class);
+            Mockito.verify(editLog).logOpRoutineLoadJob(operationCaptor.capture());
+            pauseOperation = operationCaptor.getValue();
+        }
+        ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+        try (DataOutputStream out = new DataOutputStream(bytes)) {
+            pauseOperation.write(out);
+        }
+
+        RoutineLoadOperation replayedOperation;
+        try (DataInputStream in = new DataInputStream(new ByteArrayInputStream(bytes.toByteArray()))) {
+            replayedOperation = RoutineLoadOperation.read(in);
+        }
+
+        KafkaRoutineLoadJob replayedJob = new KafkaRoutineLoadJob(jobId, "kafka_routine_load_job", 1L,
+                1L, "127.0.0.1:9020", "topic1", UserIdentity.ADMIN);
+        RoutineLoadManager routineLoadManager = new RoutineLoadManager();
+        Map<Long, RoutineLoadJob> jobs = Deencapsulation.getField(routineLoadManager, "idToRoutineLoadJob");
+        jobs.put(jobId, replayedJob);
+        routineLoadManager.replayChangeRoutineLoadJob(replayedOperation);
+
+        Assert.assertEquals(RoutineLoadJob.JobState.PAUSED, replayedJob.getState());
+        Assert.assertEquals(pauseReason.toString(), replayedJob.getPauseReason().toString());
+        Assert.assertEquals(partitionOffsets,
+                ((KafkaProgress) replayedJob.getProgress()).getOffsetByPartition());
+        Assert.assertEquals(100L, replayedJob.getRoutineLoadStatistic().totalRows);
+        Assert.assertEquals(2L, replayedJob.getRoutineLoadStatistic().errorRows);
+        Assert.assertEquals(1024L, replayedJob.getRoutineLoadStatistic().receivedBytes);
+        Assert.assertEquals(currentKafkaPartitions,
+                Deencapsulation.getField(replayedJob, "currentKafkaPartitions"));
+        Assert.assertEquals(latestOffsets,
+                Deencapsulation.getField(replayedJob, "cachedPartitionWithLatestOffsets"));
+        Assert.assertEquals(13L, replayedJob.totalLag().longValue());
+        Assert.assertTrue(replayedJob.dataSourcePropertiesJsonToString()
+                .contains("\"currentKafkaPartitions\":\"1,2\""));
+        String imageJson = GsonUtils.GSON.toJson(replayedJob);
+        Assert.assertTrue(imageJson.contains("\"pg\":"));
+        Assert.assertTrue(imageJson.contains("\"js\":"));
+        Assert.assertTrue(imageJson.contains("\"ckp\":[1,2]"));
+        Assert.assertTrue(imageJson.contains("\"cplo\":{"));
+    }
+
+    @Test
+    public void testReplayLegacyPauseKeepsReconstructedOperatorMetadata() {
+        long jobId = 1L;
+        KafkaRoutineLoadJob replayedJob = new KafkaRoutineLoadJob(jobId, "kafka_routine_load_job", 1L,
+                1L, "127.0.0.1:9020", "topic1", UserIdentity.ADMIN);
+        Map<Integer, Long> partitionOffsets = Maps.newHashMap();
+        partitionOffsets.put(1, 11L);
+        Deencapsulation.setField(replayedJob, "progress", new KafkaProgress(partitionOffsets));
+        replayedJob.getRoutineLoadStatistic().totalRows = 100L;
+        Deencapsulation.setField(replayedJob, "currentKafkaPartitions", Lists.newArrayList(1));
+        Map<Integer, Long> latestOffsets = Maps.newHashMap();
+        latestOffsets.put(1, 15L);
+        Deencapsulation.setField(replayedJob, "cachedPartitionWithLatestOffsets", latestOffsets);
+
+        RoutineLoadManager routineLoadManager = new RoutineLoadManager();
+        Map<Long, RoutineLoadJob> jobs = Deencapsulation.getField(routineLoadManager, "idToRoutineLoadJob");
+        jobs.put(jobId, replayedJob);
+        ErrorReason pauseReason = new ErrorReason(InternalErrorCode.MANUAL_PAUSE_ERR, "legacy pause");
+        routineLoadManager.replayChangeRoutineLoadJob(
+                new RoutineLoadOperation(jobId, RoutineLoadJob.JobState.PAUSED, pauseReason));
+
+        Assert.assertEquals(RoutineLoadJob.JobState.PAUSED, replayedJob.getState());
+        Assert.assertEquals(partitionOffsets,
+                ((KafkaProgress) replayedJob.getProgress()).getOffsetByPartition());
+        Assert.assertEquals(100L, replayedJob.getRoutineLoadStatistic().totalRows);
+        Assert.assertEquals(Lists.newArrayList(1),
+                Deencapsulation.getField(replayedJob, "currentKafkaPartitions"));
+        Assert.assertEquals(4L, replayedJob.totalLag().longValue());
+    }
+
     @Test
     public void testUpdateProgressWarnsWhenReadCommittedTaskHasZeroRowsAndLag() throws UserException {
         KafkaRoutineLoadJob routineLoadJob = new KafkaRoutineLoadJob(1L, "kafka_routine_load_job", 1L,