src/main/java/com/googlesource/gerrit/plugins/kinesis/CheckpointResetter.java

// Copyright (C) 2021 The Android Open Source Project
//
// 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.googlesource.gerrit.plugins.kinesis;

import static com.googlesource.gerrit.plugins.kinesis.Configuration.cosumerLeaseName;
import static software.amazon.kinesis.common.InitialPositionInStream.TRIM_HORIZON;

import com.google.common.flogger.FluentLogger;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.AttributeValueUpdate;
import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException;
import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;

@Singleton
class CheckpointResetter {
  private static final FluentLogger logger = FluentLogger.forEnclosingClass();

  private static final String LEASE_KEY_ATTRIBUTE_NAME = "leaseKey";
  private static final String LEASE_CHECKPOINT_ATTRIBUTE_VAUE = "checkpoint";
  private static final Integer DYNAMODB_RESPONSE_TIMEOUT_SECS = 5;

  private final Configuration configuration;
  private final DynamoDbAsyncClient dynamoDbAsyncClient;

  @Inject
  CheckpointResetter(Configuration configuration, DynamoDbAsyncClient dynamoDbAsyncClient) {
    this.configuration = configuration;
    this.dynamoDbAsyncClient = dynamoDbAsyncClient;
  }

  public void setAllShardsToBeginning(String streamName) {
    String leaseTable = cosumerLeaseName(configuration.getApplicationName(), streamName);

    try {
      for (String shard : getAllShards(leaseTable)) {
        logger.atInfo().log("[%s - %s] Resetting checkpoint", leaseTable, shard);

        Map<String, AttributeValue> updateKey = new HashMap<>();
        updateKey.put(LEASE_KEY_ATTRIBUTE_NAME, AttributeValue.builder().s(shard).build());

        Map<String, AttributeValueUpdate> updateValues = new HashMap<>();
        updateValues.put(
            LEASE_CHECKPOINT_ATTRIBUTE_VAUE,
            AttributeValueUpdate.builder()
                .value(AttributeValue.builder().s(TRIM_HORIZON.name().toUpperCase()).build())
                .build());

        UpdateItemResponse updateItemResponse =
            dynamoDbAsyncClient
                .updateItem(
                    UpdateItemRequest.builder()
                        .tableName(leaseTable)
                        .key(updateKey)
                        .attributeUpdates(updateValues)
                        .returnValues(ReturnValue.ALL_OLD)
                        .build())
                .get(DYNAMODB_RESPONSE_TIMEOUT_SECS, TimeUnit.SECONDS);

        logger.atInfo().log(
            "[%s - %s] Successfully reset checkpoints. old value: %s",
            leaseTable, shard, updateItemResponse);
      }
    } catch (InterruptedException e) {
      logger.atWarning().log("%s resetOffset: interrupted", leaseTable);
    } catch (ExecutionException e) {
      logger.atSevere().withCause(e).log("%s resetOffset: Error", leaseTable);
    } catch (TimeoutException e) {
      logger.atSevere().withCause(e).log("%s resetOffset: Timeout", leaseTable);
    }
  }

  private Set<String> getAllShards(String leaseTable)
      throws InterruptedException, ExecutionException, TimeoutException {
    try {
      ScanRequest scanRequest =
          ScanRequest.builder()
              .tableName(leaseTable)
              .attributesToGet(LEASE_KEY_ATTRIBUTE_NAME)
              .build();

      ScanResponse scanResponse =
          dynamoDbAsyncClient
              .scan(scanRequest)
              .get(DYNAMODB_RESPONSE_TIMEOUT_SECS, TimeUnit.SECONDS);
      return scanResponse.items().stream()
          .map(i -> i.get(LEASE_KEY_ATTRIBUTE_NAME).s())
          .collect(Collectors.toSet());
    } catch (Exception e) {
      if (e.getCause() != null && e.getCause() instanceof ResourceNotFoundException) {
        logger.atWarning().log(
            "%s resetOffset: lease table does not exist, nothing to reset.", leaseTable);
        return Collections.emptySet();
      }
      throw e;
    }
  }
}