From b7b9554364377e3f441667a19b605882dcf06003 Mon Sep 17 00:00:00 2001
From: Greyson Parrelli <greyson@signal.org>
Date: Sun, 28 Oct 2018 22:13:51 -0700
Subject: [PATCH] Prevent multiple instances of the same job running
 concurrently.

There are rare corner cases where a Job could be preempted by the
JobScheduler and then be rescheduled before the preempted job finished
running. This could create weird race conditions with bad consequences.

To fix this, we create a fun locking system that prevents two jobs with
the same UUID from running at the same time.
---
 .../securesms/jobmanager/Job.java             | 31 +++++-
 .../securesms/jobmanager/WorkLockManager.java | 94 +++++++++++++++++++
 2 files changed, 121 insertions(+), 4 deletions(-)
 create mode 100644 src/org/thoughtcrime/securesms/jobmanager/WorkLockManager.java

diff --git a/src/org/thoughtcrime/securesms/jobmanager/Job.java b/src/org/thoughtcrime/securesms/jobmanager/Job.java
index 1542418d3a..4a96ce9717 100644
--- a/src/org/thoughtcrime/securesms/jobmanager/Job.java
+++ b/src/org/thoughtcrime/securesms/jobmanager/Job.java
@@ -14,6 +14,8 @@ import org.thoughtcrime.securesms.logging.Log;
 import org.thoughtcrime.securesms.service.GenericForegroundService;
 
 import java.io.Serializable;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.UUID;
 
 import androidx.work.Data;
@@ -26,6 +28,8 @@ public abstract class Job extends Worker implements Serializable {
 
   private static final String TAG = Job.class.getSimpleName();
 
+  private static final WorkLockManager WORK_LOCK_MANAGER = new WorkLockManager();
+
   static final String KEY_RETRY_COUNT            = "Job_retry_count";
   static final String KEY_RETRY_UNTIL            = "Job_retry_until";
   static final String KEY_SUBMIT_TIME            = "Job_submit_time";
@@ -46,12 +50,28 @@ public abstract class Job extends Worker implements Serializable {
     this.parameters = parameters;
   }
 
-  @NonNull
   @Override
-  public Result doWork() {
+  public @NonNull Result doWork() {
+    log("doWork()" + logSuffix());
+
+    try (WorkLockManager.WorkLock workLock = WORK_LOCK_MANAGER.acquire(getId())) {
+      Result result = workLock.getResult();
+
+      if (result == null) {
+        result = doWorkInternal();
+        workLock.setResult(result);
+      } else {
+        log("Using result from preempted run (" + result + ")." + logSuffix());
+      }
+
+      return result;
+    }
+  }
+
+  private @NonNull Result doWorkInternal() {
     Data data = getInputData();
 
-    log("doWork()" + logSuffix());
+    log("doWorkInternal()" + logSuffix());
 
     ApplicationContext.getInstance(getApplicationContext()).injectDependencies(this);
 
@@ -104,6 +124,8 @@ public abstract class Job extends Worker implements Serializable {
     if (cancelled) {
       warn("onStopped() with cancellation signal." + logSuffix());
       onCanceled();
+    } else {
+      log("onStopped()" + logSuffix());
     }
   }
 
@@ -242,6 +264,7 @@ public abstract class Job extends Worker implements Serializable {
 
   private String logSuffix() {
     long timeSinceSubmission = System.currentTimeMillis() - getInputData().getLong(KEY_SUBMIT_TIME, 0);
-    return " (Time since submission: " + timeSinceSubmission + " ms, Run attempt: " + getRunAttemptCount() + ")";
+    return " (Time since submission: " + timeSinceSubmission + " ms, Run attempt: " + getRunAttemptCount() + ", isStopped: " + isStopped() + ")";
   }
+
 }
diff --git a/src/org/thoughtcrime/securesms/jobmanager/WorkLockManager.java b/src/org/thoughtcrime/securesms/jobmanager/WorkLockManager.java
new file mode 100644
index 0000000000..6d94b6044d
--- /dev/null
+++ b/src/org/thoughtcrime/securesms/jobmanager/WorkLockManager.java
@@ -0,0 +1,94 @@
+package org.thoughtcrime.securesms.jobmanager;
+
+import android.support.annotation.NonNull;
+import android.support.annotation.Nullable;
+
+import java.io.Closeable;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.Semaphore;
+
+import androidx.work.Worker;
+
+class WorkLockManager {
+
+  private final Map<UUID, WorkLock> locks = new HashMap<>();
+
+  WorkLock acquire(@NonNull UUID uuid) {
+    WorkLock workLock;
+
+    synchronized (this) {
+      workLock = locks.get(uuid);
+
+      if (workLock == null) {
+        workLock = new WorkLock(uuid);
+        locks.put(uuid, workLock);
+      }
+
+      workLock.increment();
+    }
+
+    workLock.getLock().acquireUninterruptibly();
+
+    return workLock;
+  }
+
+  private void release(@NonNull UUID uuid) {
+    WorkLock lock;
+
+    synchronized (this) {
+      lock = locks.get(uuid);
+
+      if (lock == null) {
+        throw new IllegalStateException("Released a lock that was already removed from use.");
+      }
+
+      if (lock.decrementAndGet() == 0) {
+        locks.remove(uuid);
+      }
+    }
+
+    lock.getLock().release();
+  }
+
+  class WorkLock implements Closeable {
+
+    private final Semaphore lock;
+    private final UUID      uuid;
+
+    private Worker.Result result;
+    private int           count;
+
+    private WorkLock(@NonNull UUID uuid) {
+      this.uuid = uuid;
+      this.lock = new Semaphore(1);
+    }
+
+    private void increment() {
+      count++;
+    }
+
+    private int decrementAndGet() {
+      count--;
+      return count;
+    }
+
+    private @NonNull Semaphore getLock() {
+      return lock;
+    }
+
+    void setResult(@NonNull Worker.Result result) {
+      this.result = result;
+    }
+
+    @Nullable Worker.Result getResult() {
+      return result;
+    }
+
+    @Override
+    public void close() {
+      WorkLockManager.this.release(uuid);
+    }
+  }
+}