ABR1: AtlasBackupService

atlasdb-backup/build.gradle

@@ -0,0 +1,10 @@
+apply from: "../gradle/shared.gradle"
+
+dependencies {
+    implementation project(':atlasdb-config')
+    implementation project(':timelock-api:timelock-api-objects')
+
+    implementation 'com.palantir.safe-logging:safe-logging'
+
+    testImplementation 'org.mockito:mockito-core'
+}

atlasdb-backup/src/main/java/com/palantir/atlasdb/backup/AtlasBackupService.java

@@ -0,0 +1,87 @@
+/*
+ * (c) Copyright 2021 Palantir Technologies Inc. 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 com.palantir.atlasdb.backup;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.palantir.atlasdb.backup.api.AtlasBackupClientBlocking;
+import com.palantir.atlasdb.timelock.api.CompleteBackupRequest;
+import com.palantir.atlasdb.timelock.api.CompleteBackupResponse;
+import com.palantir.atlasdb.timelock.api.CompletedBackup;
+import com.palantir.atlasdb.timelock.api.InProgressBackupToken;
+import com.palantir.atlasdb.timelock.api.Namespace;
+import com.palantir.atlasdb.timelock.api.PrepareBackupRequest;
+import com.palantir.atlasdb.timelock.api.PrepareBackupResponse;
+import com.palantir.conjure.java.api.config.service.ServicesConfigBlock;
+import com.palantir.dialogue.clients.DialogueClients;
+import com.palantir.dialogue.clients.DialogueClients.ReloadingFactory;
+import com.palantir.refreshable.Refreshable;
+import com.palantir.tokens.auth.AuthHeader;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+
+public final class AtlasBackupService {
+    private final AuthHeader authHeader;
+    private final AtlasBackupClientBlocking atlasBackupClientBlocking;
+    private final Map<Namespace, InProgressBackupToken> storedTokens;
+
+    @VisibleForTesting
+    AtlasBackupService(AuthHeader authHeader, AtlasBackupClientBlocking atlasBackupClientBlocking) {
+        this.authHeader = authHeader;
+        this.atlasBackupClientBlocking = atlasBackupClientBlocking;
+        this.storedTokens = new ConcurrentHashMap<>();
+    }
+
+    public static AtlasBackupService create(
+            AuthHeader authHeader, Refreshable<ServicesConfigBlock> servicesConfigBlock, String serviceName) {
+        ReloadingFactory reloadingFactory = DialogueClients.create(servicesConfigBlock);
+        AtlasBackupClientBlocking atlasBackupClientBlocking =
+                reloadingFactory.get(AtlasBackupClientBlocking.class, serviceName);
+        return new AtlasBackupService(authHeader, atlasBackupClientBlocking);
+    }
+
+    public Set<Namespace> prepareBackup(Set<Namespace> namespaces) {
+        PrepareBackupRequest request = PrepareBackupRequest.of(namespaces);
+        PrepareBackupResponse response = atlasBackupClientBlocking.prepareBackup(authHeader, request);
+
+        return response.getSuccessful().stream()
+                .peek(this::storeBackupToken)
+                .map(InProgressBackupToken::getNamespace)
+                .collect(Collectors.toSet());
+    }
+
+    private void storeBackupToken(InProgressBackupToken backupToken) {
+        storedTokens.put(backupToken.getNamespace(), backupToken);
+    }
+
+    // TODO(gs): actually persist the token using a persister passed into this class.
+    //   Then we have an atlas-side implementation of the persister that conforms with the current backup story
+    public Set<Namespace> completeBackup(Set<Namespace> namespaces) {
+        Set<InProgressBackupToken> tokens = namespaces.stream()
+                .map(storedTokens::remove)
+                .filter(Objects::nonNull)
+                .collect(Collectors.toSet());
+        CompleteBackupRequest request = CompleteBackupRequest.of(tokens);
+        CompleteBackupResponse response = atlasBackupClientBlocking.completeBackup(authHeader, request);
+
+        return response.getSuccessfulBackups().stream()
+                .map(CompletedBackup::getNamespace)
+                .collect(Collectors.toSet());
+    }
+}

atlasdb-backup/src/test/java/com/palantir/atlasdb/backup/AtlasBackupServiceTest.java

@@ -0,0 +1,109 @@
+/*
+ * (c) Copyright 2021 Palantir Technologies Inc. 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 com.palantir.atlasdb.backup;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.when;
+
+import com.google.common.collect.ImmutableSet;
+import com.palantir.atlasdb.backup.api.AtlasBackupClientBlocking;
+import com.palantir.atlasdb.timelock.api.CompleteBackupRequest;
+import com.palantir.atlasdb.timelock.api.CompleteBackupResponse;
+import com.palantir.atlasdb.timelock.api.CompletedBackup;
+import com.palantir.atlasdb.timelock.api.InProgressBackupToken;
+import com.palantir.atlasdb.timelock.api.Namespace;
+import com.palantir.atlasdb.timelock.api.PrepareBackupRequest;
+import com.palantir.atlasdb.timelock.api.PrepareBackupResponse;
+import com.palantir.lock.v2.LockToken;
+import com.palantir.tokens.auth.AuthHeader;
+import java.util.Set;
+import java.util.UUID;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class AtlasBackupServiceTest {
+    private static final Namespace NAMESPACE = Namespace.of("foo");
+    private static final Namespace OTHER_NAMESPACE = Namespace.of("other");
+    private static final InProgressBackupToken IN_PROGRESS = inProgressBackupToken(NAMESPACE);
+
+    @Mock
+    private AuthHeader authHeader;
+
+    @Mock
+    private AtlasBackupClientBlocking atlasBackupClient;
+
+    private AtlasBackupService atlasBackupService;
+
+    @Before
+    public void setup() {
+        atlasBackupService = new AtlasBackupService(authHeader, atlasBackupClient);
+    }
+
+    @Test
+    public void prepareBackupReturnsSuccessfulNamespaces() {
+        when(atlasBackupClient.prepareBackup(
+                        authHeader, PrepareBackupRequest.of(ImmutableSet.of(NAMESPACE, OTHER_NAMESPACE))))
+                .thenReturn(PrepareBackupResponse.of(ImmutableSet.of(IN_PROGRESS)));
+
+        assertThat(atlasBackupService.prepareBackup(ImmutableSet.of(NAMESPACE, OTHER_NAMESPACE)))
+                .containsExactly(NAMESPACE);
+    }
+
+    @Test
+    public void completeBackupDoesNotRunUnpreparedNamespaces() {
+        when(atlasBackupClient.completeBackup(authHeader, CompleteBackupRequest.of(ImmutableSet.of())))
+                .thenReturn(CompleteBackupResponse.of(ImmutableSet.of()));
+
+        assertThat(atlasBackupService.completeBackup(ImmutableSet.of(OTHER_NAMESPACE)))
+                .isEmpty();
+    }
+
+    @Test
+    public void completeBackupReturnsSuccessfulNamespaces() {
+        InProgressBackupToken otherInProgress = inProgressBackupToken(OTHER_NAMESPACE);
+        Set<Namespace> namespaces = ImmutableSet.of(NAMESPACE, OTHER_NAMESPACE);
+
+        when(atlasBackupClient.prepareBackup(authHeader, PrepareBackupRequest.of(namespaces)))
+                .thenReturn(PrepareBackupResponse.of(ImmutableSet.of(IN_PROGRESS, otherInProgress)));
+
+        CompletedBackup completedBackup = CompletedBackup.builder()
+                .namespace(NAMESPACE)
+                .backupStartTimestamp(2L)
+                .backupEndTimestamp(3L)
+                .build();
+        when(atlasBackupClient.completeBackup(
+                        authHeader, CompleteBackupRequest.of(ImmutableSet.of(IN_PROGRESS, otherInProgress))))
+                .thenReturn(CompleteBackupResponse.of(ImmutableSet.of(completedBackup)));
+
+        atlasBackupService.prepareBackup(namespaces);
+
+        assertThat(atlasBackupService.completeBackup(namespaces)).containsExactly(NAMESPACE);
+    }
+
+    private static InProgressBackupToken inProgressBackupToken(Namespace namespace) {
+        return InProgressBackupToken.builder()
+                .namespace(namespace)
+                .immutableTimestamp(1L)
+                .backupStartTimestamp(2L)
+                .lockToken(LockToken.of(UUID.randomUUID()))
+                .build();
+    }
+}

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/debug/LocalLockTracker.java

@@ -54,7 +54,7 @@ List<TrackedLockEvent> getLocalLockHistory() {
     void logLockResponse(Set<ConjureLockDescriptor> lockDescriptors, ConjureLockResponse response) {
         TrackedLockEvent event = getTimestampedLockEventBuilder()
                 .eventType(EventType.LOCK)
-                .eventDescription(response.accept(new ConjureLockResponse.Visitor<String>() {
+                .eventDescription(response.accept(new ConjureLockResponse.Visitor<>() {
                     @Override
                     public String visitSuccessful(SuccessfulLockResponse value) {
                         return "SUCCESS - locked " + lockDescriptors + "; obtained " + value;

changelog/@unreleased/pr-5708.v2.yml

@@ -0,0 +1,7 @@
+type: feature
+feature:
+  description: Add AtlasBackupService, which encapsulates the steps of the backup
+    process where communication with Timelock and/or knowledge of AtlasDB internals
+    is required.
+  links:
+  - https://github.com/palantir/atlasdb/pull/5708

lock-api-objects/src/main/java/com/palantir/lock/v2/TimelockService.java

@@ -46,7 +46,7 @@ default boolean isInitialized() {
 
     TimestampRange getFreshTimestamps(@Safe @QueryParam("number") int numTimestampsRequested);
 
-    // TODO (jkong): Can this be deprecated? Are there users outside of Atlas transactions?
+    // TODO(gs): Deprecate this once users outside of Atlas transactions have been eliminated
     LockImmutableTimestampResponse lockImmutableTimestamp();
 
     @DoNotDelegate

settings.gradle

@@ -1,6 +1,7 @@
 rootProject.name = 'atlasdb'
 include ":atlasdb-api"
 include ":atlasdb-autobatch"
+include ":atlasdb-backup"
 include ":atlasdb-cassandra"
 include ":atlasdb-cassandra-integration-tests"
 include ":atlasdb-cassandra-multinode-tests"

timelock-agent/src/main/java/com/palantir/timelock/paxos/TimeLockAgent.java

@@ -24,6 +24,7 @@
 import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableSet;
 import com.palantir.atlasdb.AtlasDbConstants;
+import com.palantir.atlasdb.backup.AtlasBackupResource;
 import com.palantir.atlasdb.config.AuxiliaryRemotingParameters;
 import com.palantir.atlasdb.config.ImmutableLeaderConfig;
 import com.palantir.atlasdb.config.ImmutableServerListConfig;
@@ -333,30 +334,34 @@ private void createAndRegisterResources() {
         Function<String, LockService> lockServiceGetter =
                 namespace -> namespaces.get(namespace).getLockService();
 
+        RedirectRetryTargeter redirectRetryTargeter = redirectRetryTargeter();
         if (undertowRegistrar.isPresent()) {
             Consumer<UndertowService> presentUndertowRegistrar = undertowRegistrar.get();
             registerCorruptionHandlerWrappedService(
                     presentUndertowRegistrar,
-                    ConjureTimelockResource.undertow(redirectRetryTargeter(), asyncTimelockServiceGetter));
+                    ConjureTimelockResource.undertow(redirectRetryTargeter, asyncTimelockServiceGetter));
             registerCorruptionHandlerWrappedService(
                     presentUndertowRegistrar,
-                    ConjureLockWatchingResource.undertow(redirectRetryTargeter(), asyncTimelockServiceGetter));
+                    ConjureLockWatchingResource.undertow(redirectRetryTargeter, asyncTimelockServiceGetter));
             registerCorruptionHandlerWrappedService(
-                    presentUndertowRegistrar,
-                    ConjureLockV1Resource.undertow(redirectRetryTargeter(), lockServiceGetter));
+                    presentUndertowRegistrar, ConjureLockV1Resource.undertow(redirectRetryTargeter, lockServiceGetter));
             registerCorruptionHandlerWrappedService(
                     presentUndertowRegistrar,
                     TimeLockPaxosHistoryProviderResource.undertow(corruptionComponents.localHistoryLoader()));
             registerCorruptionHandlerWrappedService(
                     presentUndertowRegistrar,
-                    MultiClientConjureTimelockResource.undertow(redirectRetryTargeter(), asyncTimelockServiceGetter));
+                    MultiClientConjureTimelockResource.undertow(redirectRetryTargeter, asyncTimelockServiceGetter));
+            registerCorruptionHandlerWrappedService(
+                    presentUndertowRegistrar,
+                    AtlasBackupResource.undertow(redirectRetryTargeter, asyncTimelockServiceGetter));
         } else {
-            registrar.accept(ConjureTimelockResource.jersey(redirectRetryTargeter(), asyncTimelockServiceGetter));
-            registrar.accept(ConjureLockWatchingResource.jersey(redirectRetryTargeter(), asyncTimelockServiceGetter));
-            registrar.accept(ConjureLockV1Resource.jersey(redirectRetryTargeter(), lockServiceGetter));
+            registrar.accept(ConjureTimelockResource.jersey(redirectRetryTargeter, asyncTimelockServiceGetter));
+            registrar.accept(ConjureLockWatchingResource.jersey(redirectRetryTargeter, asyncTimelockServiceGetter));
+            registrar.accept(ConjureLockV1Resource.jersey(redirectRetryTargeter, lockServiceGetter));
             registrar.accept(TimeLockPaxosHistoryProviderResource.jersey(corruptionComponents.localHistoryLoader()));
             registrar.accept(
-                    MultiClientConjureTimelockResource.jersey(redirectRetryTargeter(), asyncTimelockServiceGetter));
+                    MultiClientConjureTimelockResource.jersey(redirectRetryTargeter, asyncTimelockServiceGetter));
+            registrar.accept(AtlasBackupResource.jersey(redirectRetryTargeter, asyncTimelockServiceGetter));
         }
     }
 

timelock-api/src/main/conjure/timelock-api.yml

@@ -20,6 +20,10 @@ types:
       base-type: any
       external:
         java: com.palantir.lock.v2.Lease
+    LockToken:
+      base-type: any
+      external:
+        java: com.palantir.lock.v2.LockToken
     LockWatchStateUpdate:
       base-type: any
       external:
@@ -82,6 +86,29 @@ types:
         union:
           successful: SuccessfulLockResponse
           unsuccessful: UnsuccessfulLockResponse
+      InProgressBackupToken:
+        fields:
+          namespace: Namespace
+          lockToken: LockToken
+          immutableTimestamp: Long
+          backupStartTimestamp: Long
+      CompletedBackup:
+        fields:
+          namespace: Namespace
+          backupStartTimestamp: Long
+          backupEndTimestamp: Long
+      PrepareBackupRequest:
+        fields:
+          namespaces: set<Namespace>
+      PrepareBackupResponse:
+        fields:
+          successful: set<InProgressBackupToken>
+      CompleteBackupRequest:
+        fields:
+          backupTokens: set<InProgressBackupToken>
+      CompleteBackupResponse:
+        fields:
+          successfulBackups: set<CompletedBackup>
       ConjureWaitForLocksResponse:
         fields:
           wasSuccessful: boolean
@@ -117,6 +144,22 @@ types:
           leaderTimes: map<Namespace, LeaderTime>
 
 services:
+  AtlasBackupClient:
+    name: Internal backup service
+    default-auth: header
+    package: com.palantir.atlasdb.backup.api
+    base-path: /backup
+    endpoints:
+      prepareBackup:
+        http: POST /prepare
+        args:
+          request: PrepareBackupRequest
+        returns: PrepareBackupResponse
+      completeBackup:
+        http: POST /complete
+        args:
+          request: CompleteBackupRequest
+        returns: CompleteBackupResponse
   ConjureTimelockService:
     name: Timelock service
     default-auth: header