Implement primitive to run through worker's queue (#3078)

backend.native/tests/runtime/workers/worker11.kt

@@ -14,10 +14,9 @@ data class Job(val index: Int, var input: Int, var counter: Int)
 
 fun initJobs(count: Int) = Array<Job?>(count) { i -> Job(i, i * 2, i)}
 
-@Test fun runTest() {
+@Test fun runTest0() {
-    val COUNT = 100
+    val workers = Array(100, { _ -> Worker.start() })
-    val workers = Array(COUNT, { _ -> Worker.start() })
+    val jobs = initJobs(workers.size)
-    val jobs = initJobs(COUNT)
     val futures = Array(workers.size, { workerIndex ->
         workers[workerIndex].execute(TransferMode.SAFE, {
             val job = jobs[workerIndex]
@@ -40,9 +39,79 @@ fun initJobs(count: Int) = Array<Job?>(count) { i -> Job(i, i * 2, i)}
             consumed++
         }
     }
-    assertEquals(consumed, COUNT)
+    assertEquals(consumed, workers.size)
     workers.forEach {
         it.requestTermination().result
     }
     println("OK")
-}
+}
+
+val COUNT = 2
+
+@SharedImmutable
+val counters = Array(COUNT) { AtomicInt(0) }
+
+@Test fun runTest1() {
+  val workers = Array(COUNT) { Worker.start() }
+  // Ensure processQueue() can only be called on current Worker.
+  workers.forEach {
+      assertFailsWith<IllegalStateException> {
+          it.processQueue()
+      }
+  }
+  val futures = Array(workers.size) { workerIndex ->
+      workers[workerIndex].execute(TransferMode.SAFE, {
+          workerIndex
+      }) {
+          index ->
+          assertEquals(0, counters[index].value)
+          // Process following request.
+          while (!Worker.current!!.processQueue()) {}
+          // Ensure it has an effect.
+          assertEquals(1, counters[index].value)
+          // No more non-terminating tasks in this worker queue.
+          assertEquals(false, Worker.current!!.processQueue())
+      }
+  }
+  val futures2 = Array(workers.size) { workerIndex ->
+      workers[workerIndex].execute(TransferMode.SAFE, {
+          workerIndex
+      }) { index ->
+          assertEquals(0, counters[index].value)
+          counters[index].increment()
+      }
+  }
+  futures2.forEach { it.result }
+  futures.forEach { it.result }
+  workers.forEach {
+      it.requestTermination().result
+  }
+
+  // Ensure terminated workers are no longer there.
+  workers.forEach {
+    assertFailsWith<IllegalStateException> { it.execute(TransferMode.SAFE, { Unit }) { println("ERROR") } }
+  }
+}
+
+@Test fun runTest2() {
+    val workers = Array(COUNT) { Worker.start() }
+    val futures = Array(workers.size) { workerIndex ->
+        workers[workerIndex].execute(TransferMode.SAFE, { null }) {
+            // Here we processed termination request.
+            assertEquals(false, Worker.current!!.processQueue())
+        }
+    }
+    workers.forEach {
+        it.executeAfter(1000*1000, {
+            println("DELAY EXECUTED")
+            assert(false)
+        }.freeze())
+    }
+    workers.forEach {
+        it.requestTermination(processScheduledJobs = false).result
+    }
+    // Process futures, ignoring possible cancelled ones.
+    futures.forEach {
+        try { it.result } catch (e: IllegalStateException) {}
+    }
+}

runtime/src/main/cpp/Worker.cpp

@@ -60,9 +60,10 @@ enum {
 };
 
 enum JobKind {
+  JOB_NONE = 0,
   JOB_REGULAR = 1,
-  JOB_TERMINATE,
+  JOB_TERMINATE = 2,
-  JOB_EXECUTE_AFTER
+  JOB_EXECUTE_AFTER = 3
 };
 
 THREAD_LOCAL_VARIABLE KInt g_currentWorkerId = 0;
@@ -178,7 +179,7 @@ typedef KStdOrderedSet<Job, JobCompare> DelayedJobSet;
 
 class Worker {
  public:
-  Worker(KInt id, bool errorReporting) : id_(id), errorReporting_(errorReporting) {
+  Worker(KInt id, bool errorReporting) : id_(id), errorReporting_(errorReporting), terminated_(false) {
     pthread_mutex_init(&lock_, nullptr);
     pthread_cond_init(&cond_, nullptr);
   }
@@ -201,6 +202,10 @@ class Worker {
           job.terminationRequest.future->cancelUnlocked();
           break;
         }
+        case JOB_NONE: {
+          RuntimeCheck(false, "Cannot be in queue");
+          break;
+        }
       }
     }
 
@@ -228,15 +233,18 @@ class Worker {
     pthread_cond_signal(&cond_);
   }
 
-  bool waitDelayed() {
+  bool waitDelayed(bool blocking) {
     Locker locker(&lock_);
     if (delayed_.size() == 0) return false;
-    waitForQueueLocked();
+    if (blocking)
+        waitForQueueLocked();
     return true;
   }
 
-  Job getJob() {
+  Job getJob(bool blocking) {
     Locker locker(&lock_);
+    RuntimeAssert(!terminated_, "Must not be terminated");
+    if (queue_.size() == 0 && !blocking) return Job { .kind = JOB_NONE };
     waitForQueueLocked();
     auto result = queue_.front();
     queue_.pop_front();
@@ -278,6 +286,8 @@ class Worker {
     }
   }
 
+  JobKind processQueueElement(bool blocking);
+
   KInt id() const { return id_; }
 
   bool errorReporting() const { return errorReporting_; }
@@ -291,6 +301,7 @@ class Worker {
   pthread_cond_t cond_;
   // If errors to be reported on console.
   bool errorReporting_;
+  bool terminated_;
 };
 
 class State {
@@ -375,6 +386,21 @@ class State {
     return true;
   }
 
+  // Returns `true` if something was indeed processed.
+  bool processQueueUnlocked(KInt id) {
+    // Can only process queue of the current worker.
+    if (id != g_currentWorkerId) ThrowWorkerInvalidState();
+    Worker* worker = nullptr;
+    {
+      Locker locker(&lock_);
+      auto it = workers_.find(id);
+      if (it == workers_.end()) return false;
+      worker = it->second;
+    }
+    JobKind kind = worker->processQueueElement(false);
+    return kind != JOB_NONE && kind != JOB_TERMINATE;
+  }
+
   KInt stateOfFutureUnlocked(KInt id) {
     Locker locker(&lock_);
     auto it = futures_.find(id);
@@ -494,42 +520,41 @@ void Future::cancelUnlocked() {
 // Defined in RuntimeUtils.kt.
 extern "C" void ReportUnhandledException(KRef e);
 
-void* workerRoutine(void* argument) {
+JobKind Worker::processQueueElement(bool blocking) {
-  Worker* worker = reinterpret_cast<Worker*>(argument);
+  ObjHolder argumentHolder;
-
+  ObjHolder resultHolder;
-  g_currentWorkerId = worker->id();
+  if (terminated_) return JOB_TERMINATE;
-  Kotlin_initRuntimeIfNeeded();
+  Job job = getJob(blocking);
-
+  switch (job.kind) {
-  {
+    case JOB_NONE: {
-    ObjHolder argumentHolder;
+      break;
-    ObjHolder resultHolder;
+    }
-    while (true) {
+    case JOB_TERMINATE: {
-      Job job = worker->getJob();
+      if (job.terminationRequest.waitDelayed) {
-      if (job.kind == JOB_TERMINATE) {
+        if (waitDelayed(blocking)) {
-        if (job.terminationRequest.waitDelayed) {
+          putJob(job, false);
-          if (worker->waitDelayed()) {
+          return JOB_NONE;
-            worker->putJob(job, false);
-            continue;
-          }
         }
-        // Termination request, notify the future.
-        job.terminationRequest.future->storeResultUnlocked(nullptr, true);
-        theState()->removeWorkerUnlocked(worker->id());
-        break;
       }
-      if (job.kind == JOB_EXECUTE_AFTER) {
+      terminated_ = true;
-         ObjHolder operationHolder, dummyHolder;
+      // Termination request, remove the worker and notify the future.
-         KRef obj = DerefStablePointer(job.executeAfter.operation, operationHolder.slot());
+      theState()->removeWorkerUnlocked(id());
-         try {
+      job.terminationRequest.future->storeResultUnlocked(nullptr, true);
-           WorkerLaunchpad(obj, dummyHolder.slot());
+      break;
-         } catch (ExceptionObjHolder& e) {
+    }
-           if (worker->errorReporting())
+    case JOB_EXECUTE_AFTER: {
-             ReportUnhandledException(e.obj());
+      ObjHolder operationHolder, dummyHolder;
-         }
+      KRef obj = DerefStablePointer(job.executeAfter.operation, operationHolder.slot());
-         DisposeStablePointer(job.executeAfter.operation);
+      try {
-         continue;
+        WorkerLaunchpad(obj, dummyHolder.slot());
+      } catch (ExceptionObjHolder& e) {
+        if (errorReporting())
+          ReportUnhandledException(e.obj());
       }
-      RuntimeAssert(job.kind == JOB_REGULAR, "Must be regular job");
+      DisposeStablePointer(job.executeAfter.operation);
+      break;
+    }
+    case JOB_REGULAR: {
       KRef argument = AdoptStablePointer(job.regularJob.argument, argumentHolder.slot());
       KNativePtr result = nullptr;
       bool ok = true;
@@ -538,20 +563,36 @@ void* workerRoutine(void* argument) {
         argumentHolder.clear();
         // Transfer the result.
         result = transfer(&resultHolder, job.regularJob.transferMode);
-      } catch (ExceptionObjHolder& e) {
+       } catch (ExceptionObjHolder& e) {
-        ok = false;
+         ok = false;
-        if (worker->errorReporting())
+         if (errorReporting())
-          ReportUnhandledException(e.obj());
+           ReportUnhandledException(e.obj());
-      }
+       }
-      // Notify the future.
+       // Notify the future.
-      job.regularJob.future->storeResultUnlocked(result, ok);
+       job.regularJob.future->storeResultUnlocked(result, ok);
+       break;
+    }
+    default: {
+      RuntimeCheck(false, "Must be exhaustive");
     }
   }
+  return job.kind;
+}
 
-  Kotlin_deinitRuntimeIfNeeded();
+void* workerRoutine(void* argument) {
+  Worker* worker = reinterpret_cast<Worker*>(argument);
+
+  g_currentWorkerId = worker->id();
+  Kotlin_initRuntimeIfNeeded();
+
+  do {
+    if (worker->processQueueElement(true) == JOB_TERMINATE) break;
+  } while (true);
 
   konanDestructInstance(worker);
 
+  Kotlin_deinitRuntimeIfNeeded();
+
   return nullptr;
 }
 
@@ -582,6 +623,10 @@ void executeAfter(KInt id, KRef job, KLong afterMicroseconds) {
     ThrowWorkerInvalidState();
 }
 
+KBoolean processQueue(KInt id) {
+   return theState()->processQueueUnlocked(id);
+}
+
 KInt stateOfFuture(KInt id) {
   return theState()->stateOfFutureUnlocked(id);
 }
@@ -640,6 +685,10 @@ void executeAfter(KInt id, KRef job, KLong afterMicroseconds) {
   ThrowWorkerUnsupported();
 }
 
+KBoolean processQueue(KInt id) {
+  ThrowWorkerUnsupported();
+}
+
 KInt currentWorker() {
   ThrowWorkerUnsupported();
   return 0;
@@ -701,6 +750,10 @@ void Kotlin_Worker_executeAfterInternal(KInt id, KRef job, KLong afterMicrosecon
   executeAfter(id, job, afterMicroseconds);
 }
 
+KBoolean Kotlin_Worker_processQueueInternal(KInt id) {
+  return processQueue(id);
+}
+
 KInt Kotlin_Worker_stateOfFuture(KInt id) {
   return stateOfFuture(id);
 }

runtime/src/main/kotlin/kotlin/native/concurrent/Internal.kt

@@ -45,6 +45,9 @@ external internal fun executeInternal(
 @SymbolName("Kotlin_Worker_executeAfterInternal")
 external internal fun executeAfterInternal(id: Int, operation: () -> Unit, afterMicroseconds: Long): Unit
 
+@SymbolName("Kotlin_Worker_processQueueInternal")
+external internal fun processQueueInternal(id: Int): Boolean
+
 @ExportForCppRuntime
 internal fun ThrowWorkerUnsupported(): Unit =
         throw UnsupportedOperationException("Workers are not supported")

runtime/src/main/kotlin/kotlin/native/concurrent/Worker.kt

@@ -103,6 +103,18 @@ public inline class Worker @PublishedApi internal constructor(val id: Int) {
         executeAfterInternal(id, operation, afterMicroseconds)
     }
 
+
+    /**
+     * Process pending job(s) on the queue of this worker, returns `true` if something was processed
+     * and `false` otherwise. Note that jobs scheduled with [executeAfter] using non-zero timeout are
+     * not processed this way. If termination request arrives while processing the queue via this API,
+     * worker is marked as terminated and will exit once the current request is done with.
+     */
+    public fun processQueue(): Boolean = processQueueInternal(id)
+
+    /**
+     * String representation of this worker.
+     */
     override public fun toString(): String = "worker $id"
 
     /**