Reduce contention when updating `device.lastSeen`

service/src/main/java/org/whispersystems/textsecuregcm/auth/BaseAccountAuthenticator.java

@@ -83,7 +83,8 @@ public class BaseAccountAuthenticator {
       if (device.get().getAuthenticationCredentials().verify(basicCredentials.getPassword())) {
         succeeded = true;
         final Account authenticatedAccount = updateLastSeen(account.get(), device.get());
-        authenticatedAccount.setAuthenticatedDevice(device.get());
+        // the device in scope might be stale after the update, so get the latest from the authenticated account
+        authenticatedAccount.setAuthenticatedDevice(authenticatedAccount.getDevice(device.get().getId()).orElseThrow());
         return Optional.of(authenticatedAccount);
       }
 
@@ -117,7 +118,7 @@ public class BaseAccountAuthenticator {
       Metrics.summary(DAYS_SINCE_LAST_SEEN_DISTRIBUTION_NAME, IS_PRIMARY_DEVICE_TAG, String.valueOf(device.isMaster()))
           .record(Duration.ofMillis(todayInMillisWithOffset - device.getLastSeen()).toDays());
 
-      return accountsManager.updateDevice(account, device.getId(), d -> d.setLastSeen(Util.todayInMillis(clock)));
+      return accountsManager.updateDeviceLastSeen(account, device, Util.todayInMillis(clock));
     }
 
     return account;

service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java

@@ -29,6 +29,7 @@ import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import net.logstash.logback.argument.StructuredArguments;
@@ -250,12 +251,51 @@ public class AccountsManager {
 
   public Account update(Account account, Consumer<Account> updater) {
 
+    return update(account, a -> {
+      updater.accept(a);
+      // assume that all updaters passed to the public method actually modify the account
+      return true;
+    });
+  }
+
+  /**
+   * Specialized version of {@link #updateDevice(Account, long, Consumer)} that minimizes potentially contentious and
+   * redundant updates of {@code device.lastSeen}
+   */
+  public Account updateDeviceLastSeen(Account account, Device device, final long lastSeen) {
+
+    return update(account, a -> {
+
+      final Optional<Device> maybeDevice = a.getDevice(device.getId());
+
+      return maybeDevice.map(d -> {
+        if (d.getLastSeen() >= lastSeen) {
+          return false;
+        }
+
+        d.setLastSeen(lastSeen);
+
+        return true;
+
+      }).orElse(false);
+    });
+  }
+
+  /**
+   * @param account account to update
+   * @param updater must return {@code true} if the account was actually updated
+   */
+  private Account update(Account account, Function<Account, Boolean> updater) {
+
     final boolean wasVisibleBeforeUpdate = account.shouldBeVisibleInDirectory();
 
     final Account updatedAccount;
 
     try (Timer.Context ignored = updateTimer.time()) {
-      updater.accept(account);
+
+      if (!updater.apply(account)) {
+        return account;
+      }
 
       {
         // optimistically increment version
@@ -274,7 +314,11 @@ public class AccountsManager {
 
               final Optional<Account> dynamoAccount = dynamoGet(uuid);
               if (dynamoAccount.isPresent()) {
-                updater.accept(dynamoAccount.get());
+
+                if (!updater.apply(dynamoAccount.get())) {
+                  return dynamoAccount;
+                }
+
                 Account dynamoUpdatedAccount = updateWithRetries(dynamoAccount.get(),
                     updater,
                     this::dynamoUpdate,
@@ -302,7 +346,8 @@ public class AccountsManager {
     return updatedAccount;
   }
 
-  private Account updateWithRetries(Account account, Consumer<Account> updater, Consumer<Account> persister, Supplier<Account> retriever) {
+  private Account updateWithRetries(Account account, Function<Account, Boolean> updater, Consumer<Account> persister,
+      Supplier<Account> retriever) {
 
     final int maxTries = 10;
     int tries = 0;
@@ -327,7 +372,10 @@ public class AccountsManager {
       } catch (final ContestedOptimisticLockException e) {
         tries++;
         account = retriever.get();
-        updater.accept(account);
+
+        if (!updater.apply(account)) {
+          return account;
+        }
       }
     }
 
@@ -335,7 +383,11 @@ public class AccountsManager {
   }
 
   public Account updateDevice(Account account, long deviceId, Consumer<Device> deviceUpdater) {
-    return update(account, a -> a.getDevice(deviceId).ifPresent(deviceUpdater));
+    return update(account, a -> {
+      a.getDevice(deviceId).ifPresent(deviceUpdater);
+      // assume that all updaters passed to the public method actually modify the device
+      return true;
+    });
   }
 
   public Optional<Account> get(AmbiguousIdentifier identifier) {

service/src/test/java/org/whispersystems/textsecuregcm/tests/auth/BaseAccountAuthenticatorTest.java

@@ -44,13 +44,13 @@ class BaseAccountAuthenticatorTest {
 
   @BeforeEach
   void setup() {
-    accountsManager            = mock(AccountsManager.class);
+    accountsManager = mock(AccountsManager.class);
-    clock                      = mock(Clock.class);
+    clock = mock(Clock.class);
-    baseAccountAuthenticator   = new BaseAccountAuthenticator(accountsManager, clock);
+    baseAccountAuthenticator = new BaseAccountAuthenticator(accountsManager, clock);
 
-    acct1      = new Account("+14088675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
+    acct1 = new Account("+14088675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
         null, null, null, false, 0, null, yesterday, 0, null, 0, null)), null);
-    acct2      = new Account("+14098675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
+    acct2 = new Account("+14098675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
         null, null, null, false, 0, null, yesterday, 0, null, 0, null)), null);
     oldAccount = new Account("+14108675309", AuthHelper.getRandomUUID(random), Set.of(new Device(1, null, null, null,
         null, null, null, false, 0, null, oldTime, 0, null, 0, null)), null);
@@ -68,8 +68,8 @@ class BaseAccountAuthenticatorTest {
     final Account updatedAcct1 = baseAccountAuthenticator.updateLastSeen(acct1, device1);
     final Account updatedAcct2 = baseAccountAuthenticator.updateLastSeen(acct2, device2);
 
-    verify(accountsManager, never()).updateDevice(eq(acct1), anyLong(), any());
+    verify(accountsManager, never()).updateDeviceLastSeen(eq(acct1), any(), anyLong());
-    verify(accountsManager).updateDevice(eq(acct2), anyLong(), any());
+    verify(accountsManager).updateDeviceLastSeen(eq(acct2), eq(device2), anyLong());
 
     assertThat(device1.getLastSeen()).isEqualTo(yesterday);
     assertThat(device2.getLastSeen()).isEqualTo(today);
@@ -88,8 +88,8 @@ class BaseAccountAuthenticatorTest {
     final Account updatedAcct1 = baseAccountAuthenticator.updateLastSeen(acct1, device1);
     final Account updatedAcct2 = baseAccountAuthenticator.updateLastSeen(acct2, device2);
 
-    verify(accountsManager, never()).updateDevice(eq(acct1), anyLong(), any());
+    verify(accountsManager, never()).updateDeviceLastSeen(eq(acct1), any(), anyLong());
-    verify(accountsManager, never()).updateDevice(eq(acct2), anyLong(), any());
+    verify(accountsManager, never()).updateDeviceLastSeen(eq(acct2), any(), anyLong());
 
     assertThat(device1.getLastSeen()).isEqualTo(yesterday);
     assertThat(device2.getLastSeen()).isEqualTo(yesterday);
@@ -108,8 +108,8 @@ class BaseAccountAuthenticatorTest {
     final Account updatedAcct1 = baseAccountAuthenticator.updateLastSeen(acct1, device1);
     final Account updatedAcct2 = baseAccountAuthenticator.updateLastSeen(acct2, device2);
 
-    verify(accountsManager).updateDevice(eq(acct1), anyLong(), any());
+    verify(accountsManager).updateDeviceLastSeen(eq(acct1), eq(device1), anyLong());
-    verify(accountsManager).updateDevice(eq(acct2), anyLong(), any());
+    verify(accountsManager).updateDeviceLastSeen(eq(acct2), eq(device2), anyLong());
 
     assertThat(device1.getLastSeen()).isEqualTo(today);
     assertThat(device2.getLastSeen()).isEqualTo(today);
@@ -126,7 +126,7 @@ class BaseAccountAuthenticatorTest {
 
     baseAccountAuthenticator.updateLastSeen(oldAccount, device);
 
-    verify(accountsManager).updateDevice(eq(oldAccount), anyLong(), any());
+    verify(accountsManager).updateDeviceLastSeen(eq(oldAccount), eq(device), anyLong());
 
     assertThat(device.getLastSeen()).isEqualTo(today);
   }

service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java

@@ -641,4 +641,28 @@ class AccountsManagerTest {
         Arguments.of(false, true, true),
         Arguments.of(true, false, true));
   }
+
+  @ParameterizedTest
+  @MethodSource
+  void testUpdateDeviceLastSeen(final boolean expectUpdate, final long initialLastSeen, final long updatedLastSeen) {
+    final Account account = new Account("+14152222222", UUID.randomUUID(), new HashSet<>(), new byte[16]);
+    final Device device = new Device(Device.MASTER_ID, "device", "token", "salt", null, null, null, true, 1,
+        new SignedPreKey(1, "key", "sig"), initialLastSeen, 0,
+        "OWT", 0, new DeviceCapabilities());
+    account.addDevice(device);
+
+    accountsManager.updateDeviceLastSeen(account, device, updatedLastSeen);
+
+    assertEquals(expectUpdate ? updatedLastSeen : initialLastSeen, device.getLastSeen());
+    verify(accounts, expectUpdate ? times(1) : never()).update(account);
+  }
+
+  @SuppressWarnings("unused")
+  private static Stream<Arguments> testUpdateDeviceLastSeen() {
+    return Stream.of(
+        Arguments.of(true, 1, 2),
+        Arguments.of(false, 1, 1),
+        Arguments.of(false, 2, 1)
+    );
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/AccountsHelper.java

@@ -21,6 +21,7 @@ import org.mockito.MockingDetails;
 import org.mockito.stubbing.Stubbing;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
 
 public class AccountsHelper {
@@ -51,6 +52,13 @@ public class AccountsHelper {
 
       return markStale ? copyAndMarkStale(account) : account;
     });
+
+    when(mockAccountsManager.updateDeviceLastSeen(any(), any(), anyLong())).thenAnswer(answer -> {
+
+      answer.getArgument(1, Device.class).setLastSeen(answer.getArgument(2, Long.class));
+      return mockAccountsManager.update(answer.getArgument(0, Account.class), account -> {
+      });
+    });
   }
 
   public static void setupMockGet(final AccountsManager mockAccountsManager, final Set<Account> mockAccounts) {