Check if lock owner is finished on lock acquisition

Sometimes requests "die" suddenly, without going through catch/finally blocks.
If this happens, any lock they own will remain locked until it times out (which
can take hours in some cases).

This cl implicitly unlocks any lock if the owner of the lock isn't running
anymore.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=168880938

java/google/registry/model/server/Lock.java

@@ -15,12 +15,11 @@
 package google.registry.model.server;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Throwables.throwIfUnchecked;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.isAtOrAfter;
 
+import com.google.common.base.Optional;
 import com.google.common.base.Strings;
-import com.google.common.collect.ImmutableSortedSet;
 import com.googlecode.objectify.VoidWork;
 import com.googlecode.objectify.Work;
 import com.googlecode.objectify.annotation.Entity;
@@ -28,19 +27,21 @@ import com.googlecode.objectify.annotation.Id;
 import google.registry.model.ImmutableObject;
 import google.registry.model.annotations.NotBackedUp;
 import google.registry.model.annotations.NotBackedUp.Reason;
-import google.registry.util.AppEngineTimeLimiter;
 import google.registry.util.FormattingLogger;
-import java.util.HashSet;
+import google.registry.util.RequestStatusChecker;
-import java.util.Set;
+import google.registry.util.RequestStatusCheckerImpl;
-import java.util.concurrent.Callable;
-import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 
 /**
- * A lock on some shared resource. Locks are either specific to a tld or global to the entire
+ * A lock on some shared resource.
- * system, in which case a tld of null is used.
+ *
+ * <p>Locks are either specific to a tld or global to the entire system, in which case a tld of
+ * null is used.
+ *
+ * <p>This is the "barebone" lock implementation, that requires manual locking and unlocking. For
+ * safe calls that automatically lock and unlock, see LockHandler.
  */
 @Entity
 @NotBackedUp(reason = Reason.TRANSIENT)
@@ -48,13 +49,21 @@ public class Lock extends ImmutableObject {
 
   private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
 
-  /** Fudge factor to make sure we kill threads before a lock actually expires. */
-  private static final Duration LOCK_TIMEOUT_FUDGE = Duration.standardSeconds(5);
-
   /** The name of the locked resource. */
   @Id
   String lockId;
 
+  /**
+   * Unique log ID of the request that owns this lock.
+   *
+   * <p>When that request is no longer running (is finished), the lock can be considered implicitly
+   * released.
+   *
+   * <p>See {@link RequestStatusCheckerImpl#getLogId} for details about how it's created in
+   * practice.
+   */
+  String requestLogId;
+
   /** When the lock can be considered implicitly released. */
   DateTime expirationTime;
 
@@ -65,12 +74,14 @@ public class Lock extends ImmutableObject {
   private static Lock create(
       String resourceName,
       @Nullable String tld,
+      String requestLogId,
       DateTime expirationTime) {
     checkArgument(!Strings.isNullOrEmpty(resourceName), "resourceName cannot be null or empty");
     Lock instance = new Lock();
     // Add the tld to the Lock's id so that it is unique for locks acquiring the same resource
     // across different TLDs.
     instance.lockId = makeLockId(resourceName, tld);
+    instance.requestLogId = requestLogId;
     instance.expirationTime = expirationTime;
     return instance;
   }
@@ -79,15 +90,16 @@ public class Lock extends ImmutableObject {
     return String.format("%s-%s", tld, resourceName);
   }
 
-  /** Try to acquire a lock. Returns null if it can't be acquired. */
+  /** Try to acquire a lock. Returns absent if it can't be acquired. */
-  static Lock acquire(
+  public static Optional<Lock> acquire(
       final String resourceName,
       @Nullable final String tld,
-      final Duration leaseLength) {
+      final Duration leaseLength,
+      final RequestStatusChecker requestStatusChecker) {
     // It's important to use transactNew rather than transact, because a Lock can be used to control
     // access to resources like GCS that can't be transactionally rolled back. Therefore, the lock
     // must be definitively acquired before it is used, even when called inside another transaction.
-    return ofy().transactNew(new Work<Lock>() {
+    return Optional.fromNullable(ofy().transactNew(new Work<Lock>() {
       @Override
       public Lock run() {
         String lockId = makeLockId(resourceName, tld);
@@ -95,9 +107,19 @@ public class Lock extends ImmutableObject {
 
         // Checking if an unexpired lock still exists - if so, the lock can't be acquired.
         Lock lock = ofy().load().type(Lock.class).id(lockId).now();
-        if (lock != null && !isAtOrAfter(now, lock.expirationTime)) {
+        if (lock != null) {
           logger.infofmt(
-              "Existing lock is still valid now %s (until %s) lock: %s",
+              "Loaded existing lock: %s for request: %s", lock.lockId, lock.requestLogId);
+        }
+        // TODO(b/63982642): remove check on requestLogId being null once migration is done
+        // Until then we assume missing requestLogId means the app is still running (since we have
+        // no information to the contrary)
+        if (lock != null
+            && !isAtOrAfter(now, lock.expirationTime)
+            && (lock.requestLogId == null || requestStatusChecker.isRunning(lock.requestLogId))) {
+          logger.infofmt(
+              "Existing lock by request %s is still valid now %s (until %s) lock: %s",
+              lock.requestLogId,
               now,
               lock.expirationTime,
               lockId);
@@ -106,7 +128,8 @@ public class Lock extends ImmutableObject {
 
         if (lock != null) {
           logger.infofmt(
-              "Existing lock is timed out now %s (was valid until %s) lock: %s",
+              "Existing lock by request %s is timed out now %s (was valid until %s) lock: %s",
+              lock.requestLogId,
               now,
               lock.expirationTime,
               lockId);
@@ -114,6 +137,7 @@ public class Lock extends ImmutableObject {
         Lock newLock = create(
             resourceName,
             tld,
+            requestStatusChecker.getLogId(),
             now.plus(leaseLength));
         // Locks are not parented under an EntityGroupRoot (so as to avoid write contention) and
         // don't need to be backed up.
@@ -123,11 +147,11 @@ public class Lock extends ImmutableObject {
             newLock,
             lockId);
         return newLock;
-      }});
+      }}));
   }
 
   /** Release the lock. */
-  void release() {
+  public void release() {
     // Just use the default clock because we aren't actually doing anything that will use the clock.
     ofy().transact(new VoidWork() {
       @Override
@@ -151,74 +175,4 @@ public class Lock extends ImmutableObject {
         }
       }});
   }
-
-  /**
-   * Acquire one or more locks and execute a Void {@link Callable} on a thread that will be
-   * killed if it doesn't complete before the lease expires.
-   *
-   * <p>Note that locks are specific either to a given tld or to the entire system (in which case
-   * tld should be passed as null).
-   *
-   * @return whether all locks were acquired and the callable was run.
-   */
-  public static boolean executeWithLocks(
-      final Callable<Void> callable,
-      @Nullable String tld,
-      Duration leaseLength,
-      String... lockNames) {
-    try {
-      return AppEngineTimeLimiter.create().callWithTimeout(
-          new LockingCallable(callable, Strings.emptyToNull(tld), leaseLength, lockNames),
-          leaseLength.minus(LOCK_TIMEOUT_FUDGE).getMillis(),
-          TimeUnit.MILLISECONDS,
-          true);
-    } catch (Exception e) {
-      throwIfUnchecked(e);
-      throw new RuntimeException(e);
-    }
-  }
-
-  /** A {@link Callable} that acquires and releases a lock around a delegate {@link Callable}. */
-  private static class LockingCallable implements Callable<Boolean> {
-    final Callable<Void> delegate;
-    @Nullable final String tld;
-    final Duration leaseLength;
-    final Set<String> lockNames;
-
-    LockingCallable(
-        Callable<Void> delegate,
-        String tld,
-        Duration leaseLength,
-        String... lockNames) {
-      checkArgument(leaseLength.isLongerThan(LOCK_TIMEOUT_FUDGE));
-      this.delegate = delegate;
-      this.tld = tld;
-      this.leaseLength = leaseLength;
-      // Make sure we join locks in a fixed (lexicographical) order to avoid deadlock.
-      this.lockNames = ImmutableSortedSet.copyOf(lockNames);
-    }
-
-    @Override
-    public Boolean call() throws Exception {
-      Set<Lock> acquiredLocks = new HashSet<>();
-      try {
-        for (String lockName : lockNames) {
-          Lock lock = acquire(lockName, tld, leaseLength);
-          if (lock == null) {
-            logger.infofmt("Couldn't acquire lock: %s", lockName);
-            return false;
-          }
-          logger.infofmt("Acquired lock: %s", lockName);
-          acquiredLocks.add(lock);
-        }
-        delegate.call();
-        return true;
-      } finally {
-        for (Lock lock : acquiredLocks) {
-          lock.release();
-          logger.infofmt("Released lock: %s", lock.lockId);
-        }
-      }
-    }
-  }
 }

java/google/registry/request/RequestModule.java

@@ -28,7 +28,7 @@ import google.registry.request.HttpException.BadRequestException;
 import google.registry.request.HttpException.UnsupportedMediaTypeException;
 import google.registry.request.auth.AuthResult;
 import google.registry.request.lock.LockHandler;
-import google.registry.request.lock.LockHandlerPassthrough;
+import google.registry.request.lock.LockHandlerImpl;
 import google.registry.util.RequestStatusChecker;
 import google.registry.util.RequestStatusCheckerImpl;
 import java.io.IOException;
@@ -127,7 +127,7 @@ public final class RequestModule {
   }
 
   @Provides
-  static LockHandler provideLockHandler(LockHandlerPassthrough lockHandler) {
+  static LockHandler provideLockHandler(LockHandlerImpl lockHandler) {
     return lockHandler;
   }
 

java/google/registry/request/lock/BUILD

@@ -9,8 +9,10 @@ java_library(
     srcs = glob(["*.java"]),
     deps = [
         "//java/google/registry/model",
+        "//java/google/registry/util",
         "@com_google_code_findbugs_jsr305",
         "@com_google_dagger",
+        "@com_google_guava",
         "@joda_time",
     ],
 )

java/google/registry/request/lock/LockHandlerImpl.java

@@ -0,0 +1,127 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.request.lock;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Throwables.throwIfUnchecked;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Optional;
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableSortedSet;
+import google.registry.model.server.Lock;
+import google.registry.util.AppEngineTimeLimiter;
+import google.registry.util.FormattingLogger;
+import google.registry.util.RequestStatusChecker;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import javax.inject.Inject;
+import org.joda.time.Duration;
+
+/** Implementation of {@link LockHandler} that uses the datastore lock. */
+public class LockHandlerImpl implements LockHandler {
+
+  private static final long serialVersionUID = 6551645164118637767L;
+
+  private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
+
+  /** Fudge factor to make sure we kill threads before a lock actually expires. */
+  private static final Duration LOCK_TIMEOUT_FUDGE = Duration.standardSeconds(5);
+
+  @Inject RequestStatusChecker requestStatusChecker;
+
+  @Inject public LockHandlerImpl() {}
+
+  /**
+   * Acquire one or more locks and execute a Void {@link Callable}.
+   *
+   * <p>Thread will be killed if it doesn't complete before the lease expires.
+   *
+   * <p>Note that locks are specific either to a given tld or to the entire system (in which case
+   * tld should be passed as null).
+   *
+   * @return whether all locks were acquired and the callable was run.
+   */
+  @Override
+  public boolean executeWithLocks(
+      final Callable<Void> callable,
+      @Nullable String tld,
+      Duration leaseLength,
+      String... lockNames) {
+    try {
+      return AppEngineTimeLimiter.create().callWithTimeout(
+          new LockingCallable(callable, Strings.emptyToNull(tld), leaseLength, lockNames),
+          leaseLength.minus(LOCK_TIMEOUT_FUDGE).getMillis(),
+          TimeUnit.MILLISECONDS,
+          true);
+    } catch (Exception e) {
+      throwIfUnchecked(e);
+      throw new RuntimeException(e);
+    }
+  }
+
+  /** Allows injection of mock Lock in tests. */
+  @VisibleForTesting
+  Optional<Lock> acquire(String lockName, @Nullable String tld, Duration leaseLength) {
+    return Lock.acquire(lockName, tld, leaseLength, requestStatusChecker);
+  }
+
+  /** A {@link Callable} that acquires and releases a lock around a delegate {@link Callable}. */
+  private class LockingCallable implements Callable<Boolean> {
+    final Callable<Void> delegate;
+    @Nullable final String tld;
+    final Duration leaseLength;
+    final Set<String> lockNames;
+
+    LockingCallable(
+        Callable<Void> delegate,
+        String tld,
+        Duration leaseLength,
+        String... lockNames) {
+      checkArgument(leaseLength.isLongerThan(LOCK_TIMEOUT_FUDGE));
+      this.delegate = delegate;
+      this.tld = tld;
+      this.leaseLength = leaseLength;
+      // Make sure we join locks in a fixed (lexicographical) order to avoid deadlock.
+      this.lockNames = ImmutableSortedSet.copyOf(lockNames);
+    }
+
+    @Override
+    public Boolean call() throws Exception {
+      Set<Lock> acquiredLocks = new HashSet<>();
+      try {
+        for (String lockName : lockNames) {
+          Optional<Lock> lock = acquire(lockName, tld, leaseLength);
+          if (!lock.isPresent()) {
+            logger.infofmt("Couldn't acquire lock named: %s for TLD: %s", lockName, tld);
+            return false;
+          }
+          logger.infofmt("Acquired lock: %s", lock);
+          acquiredLocks.add(lock.get());
+        }
+        delegate.call();
+        return true;
+      } finally {
+        for (Lock lock : acquiredLocks) {
+          lock.release();
+          logger.infofmt("Released lock: %s", lock);
+        }
+      }
+    }
+  }
+}

java/google/registry/request/lock/LockHandlerPassthrough.java

@@ -1,54 +0,0 @@
-// Copyright 2017 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.request.lock;
-
-import google.registry.model.server.Lock;
-import java.util.concurrent.Callable;
-import javax.annotation.Nullable;
-import javax.inject.Inject;
-import org.joda.time.Duration;
-
-/**
- * Implementation of {@link LockHandler} that uses Lock as is.
- *
- * This is a temporary implementation to help migrate from Lock to LockHandler. Once the migration
- * is complete - we will create a "proper" LockHandlerImpl class and remove this one.
- *
- * TODO(guyben):delete this class once LockHandlerImpl is done.
- */
-public class LockHandlerPassthrough implements LockHandler {
-
-  private static final long serialVersionUID = 6551645164118637767L;
-
-  @Inject public LockHandlerPassthrough() {}
-
-  /**
-   * Acquire one or more locks and execute a Void {@link Callable}.
-   *
-   * <p>Runs on a thread that will be killed if it doesn't complete before the lease expires.
-   *
-   * <p>This is a simple passthrough to {@link Lock#executeWithLocks}.
-   *
-   * @return true if all locks were acquired and the callable was run; false otherwise.
-   */
-  @Override
-  public boolean executeWithLocks(
-      final Callable<Void> callable,
-      @Nullable String tld,
-      Duration leaseLength,
-      String... lockNames) {
-    return Lock.executeWithLocks(callable, tld, leaseLength, lockNames);
-  }
-}

javatests/google/registry/model/BUILD

@@ -38,6 +38,7 @@ java_library(
         "@joda_time",
         "@junit",
         "@org_joda_money",
+        "@org_mockito_all",
     ],
 )
 

javatests/google/registry/model/schema.txt

@@ -876,6 +876,7 @@ class google.registry.model.server.KmsSecretRevision {
 }
 class google.registry.model.server.Lock {
   @Id java.lang.String lockId;
+  java.lang.String requestLogId;
   org.joda.time.DateTime expirationTime;
 }
 class google.registry.model.server.ServerSecret {

javatests/google/registry/model/server/LockTest.java

@@ -15,13 +15,18 @@
 package google.registry.model.server;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
+import com.google.common.base.Optional;
 import google.registry.model.ofy.Ofy;
 import google.registry.testing.AppEngineRule;
 import google.registry.testing.ExceptionRule;
 import google.registry.testing.FakeClock;
 import google.registry.testing.InjectRule;
+import google.registry.util.RequestStatusChecker;
 import org.joda.time.Duration;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -34,6 +39,7 @@ public class LockTest {
   private static final String RESOURCE_NAME = "foo";
   private static final Duration ONE_DAY = Duration.standardDays(1);
   private static final Duration TWO_MILLIS = Duration.millis(2);
+  private static final RequestStatusChecker requestStatusChecker = mock(RequestStatusChecker.class);
 
   @Rule
   public final AppEngineRule appEngine = AppEngineRule.builder()
@@ -46,49 +52,68 @@ public class LockTest {
   @Rule
   public final ExceptionRule thrown = new ExceptionRule();
 
+  private Optional<Lock> acquire(String tld, Duration leaseLength) {
+    return Lock.acquire(RESOURCE_NAME, tld, leaseLength, requestStatusChecker);
+  }
+
+  @Before public void setUp() {
+    when(requestStatusChecker.getLogId()).thenReturn("current-request-id");
+    when(requestStatusChecker.isRunning("current-request-id")).thenReturn(true);
+  }
+
   @Test
   public void testReleasedExplicitly() throws Exception {
-    Lock lock = Lock.acquire(RESOURCE_NAME, "", ONE_DAY);
+    Optional<Lock> lock = acquire("", ONE_DAY);
-    assertThat(lock).isNotNull();
+    assertThat(lock).isPresent();
     // We can't get it again at the same time.
-    assertThat(Lock.acquire(RESOURCE_NAME, "", ONE_DAY)).isNull();
+    assertThat(acquire("", ONE_DAY)).isAbsent();
     // But if we release it, it's available.
-    lock.release();
+    lock.get().release();
-    assertThat(Lock.acquire(RESOURCE_NAME, "", ONE_DAY)).isNotNull();
+    assertThat(acquire("", ONE_DAY)).isPresent();
   }
 
   @Test
   public void testReleasedAfterTimeout() throws Exception {
     FakeClock clock = new FakeClock();
     inject.setStaticField(Ofy.class, "clock", clock);
-    assertThat(Lock.acquire(RESOURCE_NAME, "", TWO_MILLIS)).isNotNull();
+    assertThat(acquire("", TWO_MILLIS)).isPresent();
     // We can't get it again at the same time.
-    assertThat(Lock.acquire(RESOURCE_NAME, "", TWO_MILLIS)).isNull();
+    assertThat(acquire("", TWO_MILLIS)).isAbsent();
     // A second later we still can't get the lock.
     clock.advanceOneMilli();
-    assertThat(Lock.acquire(RESOURCE_NAME, "", TWO_MILLIS)).isNull();
+    assertThat(acquire("", TWO_MILLIS)).isAbsent();
     // But two seconds later we can get it.
     clock.advanceOneMilli();
-    assertThat(Lock.acquire(RESOURCE_NAME, "", TWO_MILLIS)).isNotNull();
+    assertThat(acquire("", TWO_MILLIS)).isPresent();
+  }
+
+  @Test
+  public void testReleasedAfterRequestFinish() throws Exception {
+    assertThat(acquire("", ONE_DAY)).isPresent();
+    // We can't get it again while request is active
+    assertThat(acquire("", ONE_DAY)).isAbsent();
+    // But if request is finished, we can get it.
+    when(requestStatusChecker.isRunning("current-request-id")).thenReturn(false);
+    assertThat(acquire("", ONE_DAY)).isPresent();
   }
 
   @Test
   public void testTldsAreIndependent() throws Exception {
-    Lock lockA = Lock.acquire(RESOURCE_NAME, "a", ONE_DAY);
+    Optional<Lock> lockA = acquire("a", ONE_DAY);
-    assertThat(lockA).isNotNull();
+    assertThat(lockA).isPresent();
     // For a different tld we can still get a lock with the same name.
-    Lock lockB = Lock.acquire(RESOURCE_NAME, "b", ONE_DAY);
+    Optional<Lock> lockB = acquire("b", ONE_DAY);
-    assertThat(lockB).isNotNull();
+    assertThat(lockB).isPresent();
     // We can't get lockB again at the same time.
-    assertThat(Lock.acquire(RESOURCE_NAME, "b", ONE_DAY)).isNull();
+    assertThat(acquire("b", ONE_DAY)).isAbsent();
     // Releasing lockA has no effect on lockB (even though we are still using the "b" tld).
-    lockA.release();
+    lockA.get().release();
-    assertThat(Lock.acquire(RESOURCE_NAME, "b", ONE_DAY)).isNull();
+    assertThat(acquire("b", ONE_DAY)).isAbsent();
   }
 
   @Test
   public void testFailure_emptyResourceName() throws Exception {
     thrown.expect(IllegalArgumentException.class, "resourceName cannot be null or empty");
-    Lock.acquire("", "", TWO_MILLIS);
+    Lock.acquire("", "", TWO_MILLIS, requestStatusChecker);
   }
 }

javatests/google/registry/request/lock/BUILD

@@ -0,0 +1,36 @@
+package(
+    default_testonly = 1,
+    default_visibility = ["//java/google/registry:registry_project"],
+)
+
+licenses(["notice"])  # Apache 2.0
+
+load("//java/com/google/testing/builddefs:GenTestRules.bzl", "GenTestRules")
+
+java_library(
+    name = "lock",
+    srcs = glob(["*.java"]),
+    resources = glob(["testdata/*"]),
+    deps = [
+        "//java/google/registry/model",
+        "//java/google/registry/request/lock",
+        "//javatests/google/registry/testing",
+        "//third_party/java/objectify:objectify-v4_1",
+        "@com_google_appengine_api_1_0_sdk//:testonly",
+        "@com_google_appengine_tools_appengine_gcs_client",
+        "@com_google_appengine_tools_sdk",
+        "@com_google_code_findbugs_jsr305",
+        "@com_google_guava",
+        "@com_google_truth",
+        "@javax_servlet_api",
+        "@joda_time",
+        "@junit",
+        "@org_mockito_all",
+    ],
+)
+
+GenTestRules(
+    name = "GeneratedTestRules",
+    test_files = glob(["*Test.java"]),
+    deps = [":lock"],
+)

javatests/google/registry/request/lock/LockHandlerImplTest.java

@@ -0,0 +1,132 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.request.lock;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import com.google.common.base.Optional;
+import google.registry.model.server.Lock;
+import google.registry.testing.AppEngineRule;
+import google.registry.testing.ExceptionRule;
+import java.util.concurrent.Callable;
+import javax.annotation.Nullable;
+import org.joda.time.Duration;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link LockHandler}. */
+@RunWith(JUnit4.class)
+public final class LockHandlerImplTest {
+
+  private static final Duration ONE_DAY = Duration.standardDays(1);
+
+  @Rule
+  public final AppEngineRule appEngine = AppEngineRule.builder()
+      .build();
+
+  @Rule
+  public final ExceptionRule thrown = new ExceptionRule();
+
+  private static class CountingCallable implements Callable<Void> {
+    int numCalled = 0;
+
+    @Override
+    public Void call() {
+      numCalled += 1;
+      return null;
+    }
+  }
+
+  private static class ThrowingCallable implements Callable<Void> {
+    Exception exception;
+
+    ThrowingCallable(Exception exception) {
+      this.exception = exception;
+    }
+
+    @Override
+    public Void call() throws Exception {
+      throw exception;
+    }
+  }
+
+  private boolean executeWithLocks(Callable<Void> callable, final @Nullable Lock acquiredLock) {
+    LockHandler lockHandler = new LockHandlerImpl() {
+      private static final long serialVersionUID = 0L;
+      @Override
+      Optional<Lock> acquire(String resourceName, String tld, Duration leaseLength) {
+        assertThat(resourceName).isEqualTo("resourceName");
+        assertThat(tld).isEqualTo("tld");
+        assertThat(leaseLength).isEqualTo(ONE_DAY);
+        return Optional.fromNullable(acquiredLock);
+      }
+    };
+
+    return lockHandler.executeWithLocks(callable, "tld", ONE_DAY, "resourceName");
+  }
+
+  @Before public void setUp() {
+  }
+
+  @Test
+  public void testLockSucceeds() throws Exception {
+    Lock lock = mock(Lock.class);
+    CountingCallable countingCallable = new CountingCallable();
+    assertThat(executeWithLocks(countingCallable, lock)).isTrue();
+    assertThat(countingCallable.numCalled).isEqualTo(1);
+    verify(lock, times(1)).release();
+  }
+
+  @Test
+  public void testLockSucceeds_uncheckedException() throws Exception {
+    Lock lock = mock(Lock.class);
+    Exception expectedException = new RuntimeException("test");
+    try {
+      executeWithLocks(new ThrowingCallable(expectedException), lock);
+      fail("Expected RuntimeException");
+    } catch (RuntimeException exception) {
+      assertThat(exception).isSameAs(expectedException);
+    }
+    verify(lock, times(1)).release();
+  }
+
+  @Test
+  public void testLockSucceeds_checkedException() throws Exception {
+    Lock lock = mock(Lock.class);
+    Exception expectedException = new Exception("test");
+    try {
+      executeWithLocks(new ThrowingCallable(expectedException), lock);
+      fail("Expected RuntimeException");
+    } catch (RuntimeException exception) {
+      assertThat(exception).hasCauseThat().isSameAs(expectedException);
+    }
+    verify(lock, times(1)).release();
+  }
+
+  @Test
+  public void testLockFailed() throws Exception {
+    Lock lock = null;
+    CountingCallable countingCallable = new CountingCallable();
+    assertThat(executeWithLocks(countingCallable, lock)).isFalse();
+    assertThat(countingCallable.numCalled).isEqualTo(0);
+  }
+}