Don't leak wakelock instances

Bug 31247436

Change-Id: I7c2ad956c43233b37e53856e9321cb87b01c35da
diff --git a/core/java/android/os/PowerManager.java b/core/java/android/os/PowerManager.java
index 1aed9b3..8d4d0a5 100644
--- a/core/java/android/os/PowerManager.java
+++ b/core/java/android/os/PowerManager.java
@@ -1340,6 +1340,11 @@
         }
 
         /** @hide */
+        public String getTag() {
+            return mTag;
+        }
+
+        /** @hide */
         public void setHistoryTag(String tag) {
             mHistoryTag = tag;
         }
diff --git a/services/core/java/com/android/server/job/JobServiceContext.java b/services/core/java/com/android/server/job/JobServiceContext.java
index 31528e5..5989c38 100644
--- a/services/core/java/com/android/server/job/JobServiceContext.java
+++ b/services/core/java/com/android/server/job/JobServiceContext.java
@@ -306,10 +306,24 @@
         this.service = IJobService.Stub.asInterface(service);
         final PowerManager pm =
                 (PowerManager) mContext.getSystemService(Context.POWER_SERVICE);
-        mWakeLock = pm.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, runningJob.getTag());
-        mWakeLock.setWorkSource(new WorkSource(runningJob.getSourceUid()));
-        mWakeLock.setReferenceCounted(false);
-        mWakeLock.acquire();
+        PowerManager.WakeLock wl = pm.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK,
+                runningJob.getTag());
+        wl.setWorkSource(new WorkSource(runningJob.getSourceUid()));
+        wl.setReferenceCounted(false);
+        wl.acquire();
+        synchronized (mLock) {
+            // We use a new wakelock instance per job.  In rare cases there is a race between
+            // teardown following job completion/cancellation and new job service spin-up
+            // such that if we simply assign mWakeLock to be the new instance, we orphan
+            // the currently-live lock instead of cleanly replacing it.  Watch for this and
+            // explicitly fast-forward the release if we're in that situation.
+            if (mWakeLock != null) {
+                Slog.w(TAG, "Bound new job " + runningJob + " but live wakelock " + mWakeLock
+                        + " tag=" + mWakeLock.getTag());
+                mWakeLock.release();
+            }
+            mWakeLock = wl;
+        }
         mCallbackHandler.obtainMessage(MSG_SERVICE_BOUND).sendToTarget();
     }