Fix index obs (#548)

Netflix · Aug 29, 2023 · 867913f · 867913f
1 parent 919236d
commit 867913f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/...sine-function/src/main/java/io/mantisrx/mantis/examples/sinefunction/SineFunctionJob.java b/...sine-function/src/main/java/io/mantisrx/mantis/examples/sinefunction/SineFunctionJob.java
@@ -202,6 +202,7 @@ public Observable<Observable<Integer>> call(Context context, Index index) {
             return Observable.just(
                     Observable.interval(0, period, TimeUnit.SECONDS)
                             .map(time -> {
+                                System.out.println("total worker num: " + index.getTotalNumWorkers());
                                 if (useRandom) {
                                     return randomNumGenerator.nextInt((max - min) + 1) + min;
                                 } else {

diff --git a/mantis-runtime/src/main/java/io/mantisrx/runtime/source/Index.java b/mantis-runtime/src/main/java/io/mantisrx/runtime/source/Index.java
@@ -16,14 +16,15 @@
 
 package io.mantisrx.runtime.source;
 
+import lombok.extern.slf4j.Slf4j;
 import rx.Observable;
 import rx.subjects.BehaviorSubject;
 
-
+@Slf4j
 public class Index {
 
     private final int workerIndex;
-    private final Observable<Integer> totalNumWorkersObservable;
+    private final BehaviorSubject<Integer> totalNumWorkersObservable;
 
 
     public Index(int offset, int total) {
@@ -34,14 +35,21 @@ public Index(int offset, int total) {
 
     public Index(int offset, final Observable<Integer> totalWorkerAtStageObservable) {
         this.workerIndex = offset;
-        this.totalNumWorkersObservable = totalWorkerAtStageObservable;
+        this.totalNumWorkersObservable = BehaviorSubject.create();
+        totalWorkerAtStageObservable.subscribe(this.totalNumWorkersObservable);
     }
 
     public int getWorkerIndex() {
         return workerIndex;
     }
 
     public int getTotalNumWorkers() {
+        Integer workerNum = this.totalNumWorkersObservable.getValue();
+        if (workerNum != null) {
+            return workerNum;
+        }
+
+        log.info("totalNumWorkersObservable is not ready yet, waiting.");
         return totalNumWorkersObservable.take(1).toBlocking().first();
     }