From 022dbb606f255dc81c10d3f1f8300d918071f579 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Fri, 19 Jun 2020 12:02:40 -0400 Subject: [PATCH] Count registration lock versions when crawling the account database. --- .../auth/StoredRegistrationLock.java | 4 + .../RegistrationLockVersionCounter.java | 100 +++++++++++++ .../RegistrationLockVersionCounterTest.java | 138 ++++++++++++++++++ 3 files changed, 242 insertions(+) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounter.java create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounterTest.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/auth/StoredRegistrationLock.java b/service/src/main/java/org/whispersystems/textsecuregcm/auth/StoredRegistrationLock.java index 208aaf190..b1cc99c36 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/auth/StoredRegistrationLock.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/auth/StoredRegistrationLock.java @@ -34,6 +34,10 @@ public class StoredRegistrationLock { return registrationLock.isPresent() && registrationLockSalt.isPresent(); } + public boolean hasDeprecatedPin() { + return deprecatedPin.isPresent(); + } + public long getTimeRemaining() { return TimeUnit.DAYS.toMillis(7) - (System.currentTimeMillis() - lastSeen); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounter.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounter.java new file mode 100644 index 000000000..425ff0bd0 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounter.java @@ -0,0 +1,100 @@ +package org.whispersystems.textsecuregcm.storage; + +import com.codahale.metrics.MetricRegistry; +import io.dropwizard.metrics.MetricsFactory; +import io.dropwizard.metrics.ReporterFactory; +import io.lettuce.core.KeyValue; +import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands; +import org.whispersystems.textsecuregcm.auth.StoredRegistrationLock; +import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.stream.Collectors; + +import static com.codahale.metrics.MetricRegistry.name; + +/** + * Counts the number of accounts that have the old or new (or neither) versions of a registration lock and publishes + * the results to our metric aggregator. This class can likely be removed after a few rounds of data collection. + */ +public class RegistrationLockVersionCounter extends AccountDatabaseCrawlerListener { + + private final FaultTolerantRedisCluster redisCluster; + private final MetricsFactory metricsFactory; + + static final String REGLOCK_COUNT_KEY = "ReglockVersionCounter::reglockCount"; + static final String NO_REGLOCK_KEY = "none"; + static final String PIN_ONLY_KEY = "pinOnly"; + static final String REGLOCK_ONLY_KEY = "reglockOnly"; + static final String BOTH_KEY = "both"; + + public RegistrationLockVersionCounter(final FaultTolerantRedisCluster redisCluster, final MetricsFactory metricsFactory) { + this.redisCluster = redisCluster; + this.metricsFactory = metricsFactory; + } + + @Override + public void onCrawlStart() { + redisCluster.useWriteCluster(connection -> connection.sync().hset(REGLOCK_COUNT_KEY, Map.of( + NO_REGLOCK_KEY, "0", + PIN_ONLY_KEY, "0", + REGLOCK_ONLY_KEY, "0", + BOTH_KEY, "0"))); + } + + @Override + protected void onCrawlChunk(final Optional fromUuid, final List chunkAccounts) { + int noReglockCount = 0; + int pinOnlyCount = 0; + int reglockOnlyCount = 0; + int bothCount = 0; + + for (final Account account : chunkAccounts) { + final StoredRegistrationLock storedRegistrationLock = account.getRegistrationLock(); + + if (storedRegistrationLock.hasDeprecatedPin() && storedRegistrationLock.needsFailureCredentials()) { + bothCount++; + } else if (storedRegistrationLock.hasDeprecatedPin()) { + pinOnlyCount++; + } else if (storedRegistrationLock.needsFailureCredentials()) { + reglockOnlyCount++; + } else { + noReglockCount++; + } + } + + incrementReglockCounts(noReglockCount, pinOnlyCount, reglockOnlyCount, bothCount); + } + + private void incrementReglockCounts(final int noReglockCount, final int pinOnlyCount, final int reglockOnlyCount, final int bothCount) { + redisCluster.useWriteCluster(connection -> { + final RedisAdvancedClusterCommands commands = connection.sync(); + + commands.hincrby(REGLOCK_COUNT_KEY, NO_REGLOCK_KEY, noReglockCount); + commands.hincrby(REGLOCK_COUNT_KEY, PIN_ONLY_KEY, pinOnlyCount); + commands.hincrby(REGLOCK_COUNT_KEY, REGLOCK_ONLY_KEY, reglockOnlyCount); + commands.hincrby(REGLOCK_COUNT_KEY, BOTH_KEY, bothCount); + }); + } + + @Override + public void onCrawlEnd(final Optional fromUuid) { + final Map countsByReglockType = + redisCluster.withReadCluster(connection -> connection.sync().hmget(REGLOCK_COUNT_KEY, NO_REGLOCK_KEY, PIN_ONLY_KEY, REGLOCK_ONLY_KEY, BOTH_KEY)) + .stream() + .collect(Collectors.toMap(KeyValue::getKey, keyValue -> keyValue.hasValue() ? keyValue.map(Integer::parseInt).getValue() : 0)); + + final MetricRegistry metricRegistry = new MetricRegistry(); + + for (final Map.Entry entry : countsByReglockType.entrySet()) { + metricRegistry.gauge(name(getClass(), entry.getKey()), () -> entry::getValue); + } + + for (final ReporterFactory reporterFactory : metricsFactory.getReporters()) { + reporterFactory.build(metricRegistry).report(); + } + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounterTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounterTest.java new file mode 100644 index 000000000..8cd2fcdb5 --- /dev/null +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/RegistrationLockVersionCounterTest.java @@ -0,0 +1,138 @@ +package org.whispersystems.textsecuregcm.storage; + +import com.codahale.metrics.Gauge; +import com.codahale.metrics.MetricRegistry; +import com.codahale.metrics.ScheduledReporter; +import io.dropwizard.metrics.MetricsFactory; +import io.dropwizard.metrics.ReporterFactory; +import io.lettuce.core.KeyValue; +import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands; +import org.junit.Before; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.whispersystems.textsecuregcm.auth.StoredRegistrationLock; +import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.codahale.metrics.MetricRegistry.name; +import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class RegistrationLockVersionCounterTest { + + private RedisAdvancedClusterCommands redisCommands; + private MetricsFactory metricsFactory; + + private RegistrationLockVersionCounter registrationLockVersionCounter; + + @Before + public void setUp() { + //noinspection unchecked + redisCommands = mock(RedisAdvancedClusterCommands.class); + metricsFactory = mock(MetricsFactory.class); + + registrationLockVersionCounter = new RegistrationLockVersionCounter(RedisClusterHelper.buildMockRedisCluster(redisCommands), metricsFactory); + } + + @Test + public void testOnCrawlChunkNoReglock() { + final Account account = mock(Account.class); + final StoredRegistrationLock registrationLock = mock(StoredRegistrationLock.class); + + when(account.getRegistrationLock()).thenReturn(registrationLock); + when(registrationLock.hasDeprecatedPin()).thenReturn(false); + when(registrationLock.needsFailureCredentials()).thenReturn(false); + + registrationLockVersionCounter.onCrawlChunk(Optional.empty(), List.of(account)); + + verifyCount(1, 0, 0, 0); + } + + @Test + public void testOnCrawlChunkPinOnly() { + final Account account = mock(Account.class); + final StoredRegistrationLock registrationLock = mock(StoredRegistrationLock.class); + + when(account.getRegistrationLock()).thenReturn(registrationLock); + when(registrationLock.hasDeprecatedPin()).thenReturn(true); + when(registrationLock.needsFailureCredentials()).thenReturn(false); + + registrationLockVersionCounter.onCrawlChunk(Optional.empty(), List.of(account)); + + verifyCount(0, 1, 0, 0); + } + + @Test + public void testOnCrawlChunkReglockOnly() { + final Account account = mock(Account.class); + final StoredRegistrationLock registrationLock = mock(StoredRegistrationLock.class); + + when(account.getRegistrationLock()).thenReturn(registrationLock); + when(registrationLock.hasDeprecatedPin()).thenReturn(false); + when(registrationLock.needsFailureCredentials()).thenReturn(true); + + registrationLockVersionCounter.onCrawlChunk(Optional.empty(), List.of(account)); + + verifyCount(0, 0, 1, 0); + } + + @Test + public void testOnCrawlChunkBoth() { + final Account account = mock(Account.class); + final StoredRegistrationLock registrationLock = mock(StoredRegistrationLock.class); + + when(account.getRegistrationLock()).thenReturn(registrationLock); + when(registrationLock.hasDeprecatedPin()).thenReturn(true); + when(registrationLock.needsFailureCredentials()).thenReturn(true); + + registrationLockVersionCounter.onCrawlChunk(Optional.empty(), List.of(account)); + + verifyCount(0, 0, 0, 1); + } + + private void verifyCount(final int noReglock, final int pinOnly, final int reglockOnly, final int both) { + verify(redisCommands).hincrby(RegistrationLockVersionCounter.REGLOCK_COUNT_KEY, RegistrationLockVersionCounter.NO_REGLOCK_KEY, noReglock); + verify(redisCommands).hincrby(RegistrationLockVersionCounter.REGLOCK_COUNT_KEY, RegistrationLockVersionCounter.PIN_ONLY_KEY, pinOnly); + verify(redisCommands).hincrby(RegistrationLockVersionCounter.REGLOCK_COUNT_KEY, RegistrationLockVersionCounter.REGLOCK_ONLY_KEY, reglockOnly); + verify(redisCommands).hincrby(RegistrationLockVersionCounter.REGLOCK_COUNT_KEY, RegistrationLockVersionCounter.BOTH_KEY, both); + } + + @Test + public void testOnCrawlEnd() { + final int noReglockCount = 21; + final int pinOnlyCount = 7; + final int reglockOnlyCount = 83; + + final ReporterFactory reporterFactory = mock(ReporterFactory.class); + final ScheduledReporter reporter = mock(ScheduledReporter.class); + + when(metricsFactory.getReporters()).thenReturn(List.of(reporterFactory)); + + final ArgumentCaptor registryCaptor = ArgumentCaptor.forClass(MetricRegistry.class); + when(reporterFactory.build(any())).thenReturn(reporter); + + when(redisCommands.hmget(eq(RegistrationLockVersionCounter.REGLOCK_COUNT_KEY), any())).thenReturn(List.of( + KeyValue.just(RegistrationLockVersionCounter.NO_REGLOCK_KEY, String.valueOf(noReglockCount)), + KeyValue.just(RegistrationLockVersionCounter.PIN_ONLY_KEY, String.valueOf(pinOnlyCount)), + KeyValue.just(RegistrationLockVersionCounter.REGLOCK_ONLY_KEY, String.valueOf(reglockOnlyCount)), + KeyValue.empty(RegistrationLockVersionCounter.BOTH_KEY))); + + registrationLockVersionCounter.onCrawlEnd(Optional.empty()); + + verify(reporterFactory).build(registryCaptor.capture()); + verify(reporter).report(); + + @SuppressWarnings("rawtypes") final Map gauges = registryCaptor.getValue().getGauges(); + assertEquals(noReglockCount, gauges.get(name(RegistrationLockVersionCounter.class, RegistrationLockVersionCounter.NO_REGLOCK_KEY)).getValue()); + assertEquals(pinOnlyCount, gauges.get(name(RegistrationLockVersionCounter.class, RegistrationLockVersionCounter.PIN_ONLY_KEY)).getValue()); + assertEquals(reglockOnlyCount, gauges.get(name(RegistrationLockVersionCounter.class, RegistrationLockVersionCounter.REGLOCK_ONLY_KEY)).getValue()); + assertEquals(0, gauges.get(name(RegistrationLockVersionCounter.class, RegistrationLockVersionCounter.BOTH_KEY)).getValue()); + } +}