From 2e01da5ec1b545e13e502dc4e9b3b6cc9c24fc56 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Mon, 11 Jan 2021 18:51:20 -0500 Subject: [PATCH] Add a task to enable/disable accelerated crawling. --- .../textsecuregcm/WhisperServerService.java | 2 ++ .../storage/AccountDatabaseCrawler.java | 4 +-- .../storage/AccountDatabaseCrawlerCache.java | 8 +++-- .../workers/SetCrawlerAccelerationTask.java | 36 +++++++++++++++++++ .../storage/AccountDatabaseCrawlerTest.java | 4 +-- 5 files changed, 48 insertions(+), 6 deletions(-) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/SetCrawlerAccelerationTask.java diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java index 2d96a0f58..984284034 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java @@ -150,6 +150,7 @@ import org.whispersystems.textsecuregcm.workers.DisableRequestLoggingTask; import org.whispersystems.textsecuregcm.workers.EnableRequestLoggingTask; import org.whispersystems.textsecuregcm.workers.GetRedisCommandStatsCommand; import org.whispersystems.textsecuregcm.workers.GetRedisSlowlogCommand; +import org.whispersystems.textsecuregcm.workers.SetCrawlerAccelerationTask; import org.whispersystems.textsecuregcm.workers.VacuumCommand; import org.whispersystems.textsecuregcm.workers.ZkParamsCommand; import org.whispersystems.websocket.WebSocketResourceProviderFactory; @@ -448,6 +449,7 @@ public class WhisperServerService extends Application listener.onCrawlEnd(fromUuid)); cache.setLastUuid(Optional.empty()); - cache.clearAccelerate(); + cache.setAccelerated(false); } else { try { for (AccountDatabaseCrawlerListener listener : listeners) { @@ -131,7 +131,7 @@ public class AccountDatabaseCrawler implements Managed, Runnable { cache.setLastUuid(Optional.of(chunkAccounts.get(chunkAccounts.size() - 1).getUuid())); } catch (AccountDatabaseCrawlerRestartException e) { cache.setLastUuid(Optional.empty()); - cache.clearAccelerate(); + cache.setAccelerated(false); } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountDatabaseCrawlerCache.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountDatabaseCrawlerCache.java index c8a134f74..18740a1e9 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountDatabaseCrawlerCache.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountDatabaseCrawlerCache.java @@ -31,8 +31,12 @@ public class AccountDatabaseCrawlerCache { this.unlockClusterScript = ClusterLuaScript.fromResource(cacheCluster, "lua/account_database_crawler/unlock.lua", ScriptOutputType.INTEGER); } - public void clearAccelerate() { - cacheCluster.useCluster(connection -> connection.sync().del(ACCELERATE_KEY)); + public void setAccelerated(final boolean accelerated) { + if (accelerated) { + cacheCluster.useCluster(connection -> connection.sync().set(ACCELERATE_KEY, "1")); + } else { + cacheCluster.useCluster(connection -> connection.sync().del(ACCELERATE_KEY)); + } } public boolean isAccelerated() { diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/SetCrawlerAccelerationTask.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/SetCrawlerAccelerationTask.java new file mode 100644 index 000000000..a7cd2e560 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/SetCrawlerAccelerationTask.java @@ -0,0 +1,36 @@ +/* + * Copyright 2021 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ + +package org.whispersystems.textsecuregcm.workers; + +import io.dropwizard.servlets.tasks.Task; +import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerCache; + +import java.io.PrintWriter; +import java.util.List; +import java.util.Map; + +public class SetCrawlerAccelerationTask extends Task { + + private final AccountDatabaseCrawlerCache crawlerCache; + + public SetCrawlerAccelerationTask(final AccountDatabaseCrawlerCache crawlerCache) { + super("set-crawler-accelerated"); + + this.crawlerCache = crawlerCache; + } + + @Override + public void execute(final Map> parameters, final PrintWriter out) { + if (parameters.containsKey("accelerated") && parameters.get("accelerated").size() == 1) { + final boolean accelerated = "true".equalsIgnoreCase(parameters.get("accelerated").get(0)); + + crawlerCache.setAccelerated(accelerated); + out.println("Set accelerated: " + accelerated); + } else { + out.println("Usage: set-crawler-accelerated?accelerated=[true|false]"); + } + } +} diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountDatabaseCrawlerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountDatabaseCrawlerTest.java index de9cc1ef1..f9dfec1c0 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountDatabaseCrawlerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountDatabaseCrawlerTest.java @@ -148,7 +148,7 @@ public class AccountDatabaseCrawlerTest { verify(account2, times(0)).getNumber(); verify(listener, times(1)).timeAndProcessCrawlChunk(eq(Optional.of(ACCOUNT1)), eq(Arrays.asList(account2))); verify(cache, times(1)).setLastUuid(eq(Optional.empty())); - verify(cache, times(1)).clearAccelerate(); + verify(cache, times(1)).setAccelerated(false); verify(cache, times(1)).isAccelerated(); verify(cache, times(1)).releaseActiveWork(any(String.class)); @@ -175,7 +175,7 @@ public class AccountDatabaseCrawlerTest { verify(account2, times(0)).getNumber(); verify(listener, times(1)).onCrawlEnd(eq(Optional.of(ACCOUNT2))); verify(cache, times(1)).setLastUuid(eq(Optional.empty())); - verify(cache, times(1)).clearAccelerate(); + verify(cache, times(1)).setAccelerated(false); verify(cache, times(1)).isAccelerated(); verify(cache, times(1)).releaseActiveWork(any(String.class));