pipeline context example

src/main/java/ru/nbch/credit_tracker/model/PipelineContext.java

@@ -0,0 +1,21 @@
+package ru.nbch.credit_tracker.model;
+
+import java.util.EnumMap;
+
+public class PipelineContext<T extends Enum<T>> {
+    private final EnumMap<T, Object> storage;
+
+    public PipelineContext(Class<T> keyClass) {
+        this.storage = new EnumMap<>(keyClass);
+    }
+
+    public void store(T key, Object value) {
+        storage.put(key, value);
+    }
+
+    @SuppressWarnings("unchecked")
+    public <S> S get(T key) {
+        Object o = storage.get(key);
+        return (S) o;
+    }
+}

src/main/java/ru/nbch/credit_tracker/model/StagePayload.java

@@ -2,4 +2,6 @@ package ru.nbch.credit_tracker.model;
 
 public interface StagePayload {
     boolean isPoison();
+    default void setContext(PipelineContext<?> context) {}
+    default PipelineContext<?> getContext() {return null;}
 }

src/main/java/ru/nbch/credit_tracker/model/impl/PackageVariable.java

@@ -0,0 +1,5 @@
+package ru.nbch.credit_tracker.model.impl;
+
+public enum PackageVariable {
+    packageId,
+}

src/main/java/ru/nbch/credit_tracker/model/impl/SubjectPayload.java

@@ -2,6 +2,7 @@ package ru.nbch.credit_tracker.model.impl;
 
 import java.util.List;
 import ru.nbch.credit_tracker.enums.PoisonType;
+import ru.nbch.credit_tracker.model.PipelineContext;
 import ru.nbch.credit_tracker.model.StagePayload;
 import ru.nbch.credit_tracker.model.SubjectProfile;
 
@@ -11,6 +12,7 @@ public class SubjectPayload implements StagePayload {
     private final String memberCode;
     private final String packageName;
     private final List<SubjectProfile> subjects;
+    private PipelineContext<PackageVariable> ctx;
     public static final SubjectPayload POISON = new SubjectPayload(null, PoisonType.NORMALLY);
     private final PoisonType poison;
 
@@ -30,6 +32,16 @@ public class SubjectPayload implements StagePayload {
         this.poison = PoisonType.NONE;
     }
 
+    @Override
+    public void setContext(PipelineContext<?> context) {
+        this.ctx = (PipelineContext<PackageVariable>) context;
+    }
+
+    @Override
+    public PipelineContext<?> getContext() {
+        return this.ctx;
+    }
+
     public Long getPackageId() {
         return packageId;
     }

src/main/java/ru/nbch/credit_tracker/service/task/Handler.java

@@ -1,7 +1,17 @@
 package ru.nbch.credit_tracker.service.task;
 
-public interface Handler<X, Z> {
+import ru.nbch.credit_tracker.model.PipelineContext;
+import ru.nbch.credit_tracker.model.StagePayload;
+
+public interface Handler<X extends StagePayload, Z extends StagePayload> {
     Z handle(X taskPayload);
+
+    @SuppressWarnings("unchecked")
+    default <T extends Enum<T>> PipelineContext<T> obtainContext(Class<T> keyClass, X input) {
+        PipelineContext<?> ctx = input.getContext();
+        return (PipelineContext<T>) ctx;
+    }
+
     default void handleOnPoison(X poison){
         //do nothing
     }

src/main/java/ru/nbch/credit_tracker/service/task/PipelinePackageFabric.java

@@ -8,6 +8,7 @@ import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.stereotype.Component;
 import ru.nbch.credit_tracker.config.settings.CreditTrackerSettings;
 import ru.nbch.credit_tracker.config.task.StageQueueFactory;
+import ru.nbch.credit_tracker.model.impl.PackageVariable;
 import ru.nbch.credit_tracker.model.impl.PartFilePayload;
 import ru.nbch.credit_tracker.model.impl.SignalXmlFilePayload;
 import ru.nbch.credit_tracker.model.impl.SubjectPayload;
@@ -64,6 +65,7 @@ public class PipelinePackageFabric implements InitializingBean {
                 .fillInitialQueue(pipelineStarterFactory.apply(setupParam))
                 .pollFinishingQueue(pipelineTerminatorFactory.get())
                 .backoff(pipelineBackoff)
+                .context(PackageVariable.class)
                 .build();
     }
 

src/main/java/ru/nbch/credit_tracker/service/task/Stage.java

@@ -12,6 +12,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.exception.ExceptionUtils;
+import ru.nbch.credit_tracker.model.PipelineContext;
 import ru.nbch.credit_tracker.model.StagePayload;
 import ru.nbch.credit_tracker.utils.LogUtil;
 import ru.nbch.credit_tracker.utils.TaskTimer;
@@ -30,6 +31,7 @@ public class Stage<X extends StagePayload, Z extends StagePayload> implements Ru
 
     private static final ThreadPoolExecutor.CallerRunsPolicy callerRunsPolicy
         = new ThreadPoolExecutor.CallerRunsPolicy();
+    private PipelineContext<?> context;
 
     public Stage(Handler<X, Z> dispatcher,
                  BlockingQueue<X> in,
@@ -106,7 +108,11 @@ public class Stage<X extends StagePayload, Z extends StagePayload> implements Ru
                             log.info("stage {} task {} started.", dispatcher.getClass().getSimpleName(), finalCounter);
                             Z p;
                             try (var m = stopWatch.start(dispatcherName + "-handler")) {
+                                if (b.getContext() == null) {
+                                    b.setContext(context);
+                                }
                                 p = dispatcher.handle(b);
+                                p.setContext(b.getContext()); // нужно т.к. в конце придет в терминатор
                             }
                             if (Thread.currentThread().isInterrupted()) {
                                 log.info("after {} task {} finished, worker has been interrupted. not putting result in queue.",
@@ -210,4 +216,12 @@ public class Stage<X extends StagePayload, Z extends StagePayload> implements Ru
     public void setBackoff(PipelineBackoff pipelineBackoff) {
         this.pipelineBackoff = pipelineBackoff;
     }
+
+    public void setContext(PipelineContext<?> context) {
+        this.context = context;
+    }
+
+    public PipelineContext<?> getContext() {
+        return context;
+    }
 }

src/main/java/ru/nbch/credit_tracker/service/task/impl/FidNameEnricher.java

@@ -10,6 +10,7 @@ import ru.nbch.credit_tracker.entities.app.PhoneFidMap;
 import ru.nbch.credit_tracker.entities.indic.Name;
 import ru.nbch.credit_tracker.model.NameData;
 import ru.nbch.credit_tracker.model.SubjectProfile;
+import ru.nbch.credit_tracker.model.impl.PackageVariable;
 import ru.nbch.credit_tracker.model.impl.SubjectPayload;
 import ru.nbch.credit_tracker.service.indic.PhoneService2;
 import ru.nbch.credit_tracker.service.task.Handler;
@@ -35,7 +36,9 @@ public class FidNameEnricher implements Handler<SubjectPayload, SubjectPayload>
     @Override
     public SubjectPayload handle(SubjectPayload taskPayload) {
         List<SubjectProfile> subjects = taskPayload.getSubjects();
-        log.debug("NameEnricherStage handle by packageId: {}, size of bucket for save: {}", taskPayload.getPackageId(), subjects.size());
+        var ctx = obtainContext(PackageVariable.class, taskPayload);
+        Long packageId = ctx.get(PackageVariable.packageId);
+        log.debug("NameEnricherStage handle by packageId: {}, size of bucket for save: {}", packageId, subjects.size());
 
         if (subjects.getFirst().getFid() == null) {
             processByNumber(subjects);

src/main/java/ru/nbch/credit_tracker/service/task/impl/KickPipeLine.java

@@ -12,7 +12,9 @@ import ru.nbch.credit_tracker.enums.PoisonType;
 import ru.nbch.credit_tracker.exceptions.SignalClientException;
 import ru.nbch.credit_tracker.mapper.signal.PhoneFidMapMapper;
 import ru.nbch.credit_tracker.mapper.signal.UserPackageMapper;
+import ru.nbch.credit_tracker.model.PipelineContext;
 import ru.nbch.credit_tracker.model.SubjectProfile;
+import ru.nbch.credit_tracker.model.impl.PackageVariable;
 import ru.nbch.credit_tracker.model.impl.SubjectPayload;
 import ru.nbch.credit_tracker.service.task.SetupParam;
 import ru.nbch.credit_tracker.service.task.setup.IPipeLineStarter;
@@ -30,6 +32,7 @@ public abstract class KickPipeLine implements IPipeLineStarter<SubjectPayload> {
     private final int BATCH_SIZE;
     private boolean isInterrupted = false;
     protected int offset = 0;
+    private PipelineContext<PackageVariable> context;
 
 
     public KickPipeLine(StaxXmlProcessor xmlProcessor,
@@ -52,10 +55,17 @@ public abstract class KickPipeLine implements IPipeLineStarter<SubjectPayload> {
         }
     }
 
+    @SuppressWarnings("unchecked")
+    @Override
+    public void setContext(PipelineContext<?> context) {
+        this.context = (PipelineContext<PackageVariable>) context;
+    }
+
     @Override
     public void startFillingQueue(BlockingQueue<SubjectPayload> firstQueue) {
         try {
             log.info("Start kicker: mode:{}, packageId: {};", mode(), packageId);
+            context.store(PackageVariable.packageId, packageId);
             int limit = BATCH_SIZE;
             UserPackage userPackage = userPackageMapper.findPackage(packageId);
             List<PhoneFidMap> page = nextPage(limit);
@@ -69,7 +79,9 @@ public abstract class KickPipeLine implements IPipeLineStarter<SubjectPayload> {
                     .phone(phoneFidMap.getPhone())
                     .build()).toList();
                 try {
-                    firstQueue.put(new SubjectPayload(packageId, userPackage.getUserId(), userPackage.getPackageName(), subjectProfiles));
+                    SubjectPayload payload = new SubjectPayload(packageId, userPackage.getUserId(), userPackage.getPackageName(), subjectProfiles);
+                    payload.setContext(context);
+                    firstQueue.put(payload);
                 } catch (InterruptedException e) {
                     isInterrupted = true;
                     log.error("Pipeline starter was interrupted!");

src/main/java/ru/nbch/credit_tracker/service/task/setup/IPipeLineStarter.java

@@ -1,8 +1,11 @@
 package ru.nbch.credit_tracker.service.task.setup;
 
 import java.util.concurrent.BlockingQueue;
+import ru.nbch.credit_tracker.model.PipelineContext;
 import ru.nbch.credit_tracker.model.StagePayload;
 
 public interface IPipeLineStarter<Q extends StagePayload> {
     void startFillingQueue(BlockingQueue<Q> queue);
+
+    default void setContext(PipelineContext<?> context) {}
 }

src/main/java/ru/nbch/credit_tracker/service/task/setup/PipeLineBuilderV2.java

@@ -4,6 +4,7 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.BlockingQueue;
+import ru.nbch.credit_tracker.model.PipelineContext;
 import ru.nbch.credit_tracker.model.StagePayload;
 import ru.nbch.credit_tracker.service.task.Pipeline;
 import ru.nbch.credit_tracker.service.task.PipelineBackoff;
@@ -17,6 +18,7 @@ public class PipeLineBuilderV2<FIRST extends StagePayload, LAST extends StagePay
     private IPipeLineStarter<FIRST> pipeLineStarter;
     private IPipeLineTerminator<LAST, ?> pipeLineTerminator;
     private PipelineBackoff backoff;
+    private PipelineContext<?> context;
 
     private PipeLineBuilderV2() {
         this.stages = new ArrayList<>();
@@ -60,14 +62,22 @@ public class PipeLineBuilderV2<FIRST extends StagePayload, LAST extends StagePay
         return this;
     }
 
+    public <T extends Enum<T>>
+    PipeLineBuilderV2<FIRST, LAST> context(Class<T> contextKeyClass) {
+        this.context = new PipelineContext<>(contextKeyClass);
+        return this;
+    }
+
     public Pipeline build() {
         Pipeline p = new Pipeline();
         p.setStarter(() -> {
             stages.forEach(stage -> {
                 stage.setPipeline(p);
                 stage.setBackoff(backoff);
+                stage.setContext(context);
                 stage.startRunning();
             });
+            pipeLineStarter.setContext(context);
             pipeLineStarter.startFillingQueue(firstStage.getIn());
         });
         pipeLineTerminator.setPollingQueue(lastQueue);