org.apache.iceberg.aws.glue.DynamoLockManager Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of iceberg-aws Show documentation
A table format for huge analytic datasets
The newest version!
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.iceberg.aws.glue;

import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import org.apache.iceberg.CatalogProperties;
import org.apache.iceberg.aws.AwsClientFactories;
import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.relocated.com.google.common.collect.Maps;
import org.apache.iceberg.util.LockManagers;
import org.apache.iceberg.util.Tasks;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.BillingMode;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest;
import software.amazon.awssdk.services.dynamodb.model.DescribeTableResponse;
import software.amazon.awssdk.services.dynamodb.model.DynamoDbException;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.InternalServerErrorException;
import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
import software.amazon.awssdk.services.dynamodb.model.KeyType;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughputExceededException;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
import software.amazon.awssdk.services.dynamodb.model.RequestLimitExceededException;
import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException;
import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
import software.amazon.awssdk.services.dynamodb.model.TableStatus;
import software.amazon.awssdk.services.dynamodb.model.TransactionConflictException;

/**
 * DynamoDB implementation for the lock manager.
 */
class DynamoLockManager extends LockManagers.BaseLockManager {

  private static final Logger LOG = LoggerFactory.getLogger(DynamoLockManager.class);

  private static final String COL_LOCK_ENTITY_ID = "entityId";
  private static final String COL_LEASE_DURATION_MS = "leaseDurationMs";
  private static final String COL_VERSION = "version";
  private static final String COL_LOCK_OWNER_ID = "ownerId";

  private static final String CONDITION_LOCK_ID_MATCH = String.format(
      "%s = :eid AND %s = :oid",
      COL_LOCK_ENTITY_ID, COL_LOCK_OWNER_ID);
  private static final String CONDITION_LOCK_ENTITY_NOT_EXIST = String.format(
      "attribute_not_exists(%s)",
      COL_LOCK_ENTITY_ID);
  private static final String CONDITION_LOCK_ENTITY_NOT_EXIST_OR_VERSION_MATCH = String.format(
      "attribute_not_exists(%s) OR (%s = :eid AND %s = :vid)",
      COL_LOCK_ENTITY_ID, COL_LOCK_ENTITY_ID, COL_VERSION);

  private static final int LOCK_TABLE_CREATION_WAIT_ATTEMPTS_MAX = 5;
  private static final int RELEASE_RETRY_ATTEMPTS_MAX = 5;

  private static final List LOCK_TABLE_SCHEMA = Lists.newArrayList(
      KeySchemaElement.builder()
          .attributeName(COL_LOCK_ENTITY_ID)
          .keyType(KeyType.HASH)
          .build());

  private static final List LOCK_TABLE_COL_DEFINITIONS = Lists.newArrayList(
      AttributeDefinition.builder()
          .attributeName(COL_LOCK_ENTITY_ID)
          .attributeType(ScalarAttributeType.S)
          .build());

  private final Map heartbeats = Maps.newHashMap();

  private DynamoDbClient dynamo;
  private String lockTableName;

  /**
   * constructor for dynamic initialization, {@link #initialize(Map)} must be called later.
   */
  DynamoLockManager() {
  }

  /**
   * constructor used for testing purpose
   * @param dynamo dynamo client
   * @param lockTableName lock table name
   */
  DynamoLockManager(DynamoDbClient dynamo, String lockTableName) {
    super.initialize(Maps.newHashMap());
    this.dynamo = dynamo;
    this.lockTableName = lockTableName;
    ensureLockTableExistsOrCreate();
  }

  private void ensureLockTableExistsOrCreate() {

    if (tableExists(lockTableName)) {
      return;
    }

    LOG.info("Dynamo lock table {} not found, trying to create", lockTableName);
    dynamo.createTable(CreateTableRequest.builder()
        .tableName(lockTableName)
        .keySchema(lockTableSchema())
        .attributeDefinitions(lockTableColDefinitions())
        .billingMode(BillingMode.PAY_PER_REQUEST)
        .build());

    Tasks.foreach(lockTableName)
        .retry(LOCK_TABLE_CREATION_WAIT_ATTEMPTS_MAX)
        .throwFailureWhenFinished()
        .onlyRetryOn(IllegalStateException.class)
        .run(this::checkTableActive);
  }

  @VisibleForTesting
  boolean tableExists(String tableName) {
    try {
      dynamo.describeTable(DescribeTableRequest.builder()
          .tableName(tableName)
          .build());
      return true;
    } catch (ResourceNotFoundException e) {
      return false;
    }
  }

  private void checkTableActive(String tableName) {
    try {
      DescribeTableResponse response = dynamo.describeTable(DescribeTableRequest.builder()
          .tableName(tableName)
          .build());
      TableStatus currentStatus = response.table().tableStatus();
      if (!currentStatus.equals(TableStatus.ACTIVE)) {
        throw new IllegalStateException(String.format("Dynamo table %s is not active, current status: %s",
            tableName, currentStatus));
      }
    } catch (ResourceNotFoundException e) {
      throw new IllegalStateException(String.format("Cannot find Dynamo table %s", tableName));
    }
  }

  @Override
  public void initialize(Map properties) {
    super.initialize(properties);
    this.dynamo = AwsClientFactories.from(properties).dynamo();
    this.lockTableName = properties.get(CatalogProperties.LOCK_TABLE);
    Preconditions.checkNotNull(lockTableName, "DynamoDB lock table name must not be null");
    ensureLockTableExistsOrCreate();
  }

  @Override
  public boolean acquire(String entityId, String ownerId) {
    try {
      Tasks.foreach(entityId)
          .throwFailureWhenFinished()
          .retry(Integer.MAX_VALUE - 1)
          .exponentialBackoff(acquireIntervalMs(), acquireIntervalMs(), acquireTimeoutMs(), 1)
          .onlyRetryOn(
              ConditionalCheckFailedException.class,
              ProvisionedThroughputExceededException.class,
              TransactionConflictException.class,
              RequestLimitExceededException.class,
              InternalServerErrorException.class)
          .run(id -> acquireOnce(id, ownerId));
      return true;
    } catch (DynamoDbException e) {
      return false;
    }
  }

  @VisibleForTesting
  void acquireOnce(String entityId, String ownerId) {
    GetItemResponse response = dynamo.getItem(GetItemRequest.builder()
        .tableName(lockTableName)
        .key(toKey(entityId))
        .build());

    if (!response.hasItem()) {
      dynamo.putItem(PutItemRequest.builder()
          .tableName(lockTableName)
          .item(toNewItem(entityId, ownerId, heartbeatTimeoutMs()))
          .conditionExpression(CONDITION_LOCK_ENTITY_NOT_EXIST)
          .build());
    } else {
      Map currentItem = response.item();

      try {
        Thread.sleep(Long.parseLong(currentItem.get(COL_LEASE_DURATION_MS).n()));
      } catch (InterruptedException e) {
        throw new IllegalStateException(
            String.format("Fail to acquire lock %s by %s, interrupted during sleep", entityId, ownerId), e);
      }

      dynamo.putItem(PutItemRequest.builder()
          .tableName(lockTableName)
          .item(toNewItem(entityId, ownerId, heartbeatTimeoutMs()))
          .conditionExpression(CONDITION_LOCK_ENTITY_NOT_EXIST_OR_VERSION_MATCH)
          .expressionAttributeValues(ImmutableMap.of(
              ":eid", AttributeValue.builder().s(entityId).build(),
              ":vid", AttributeValue.builder().s(currentItem.get(COL_VERSION).s()).build()))
          .build());
    }

    startNewHeartbeat(entityId, ownerId);
  }

  private void startNewHeartbeat(String entityId, String ownerId) {
    if (heartbeats.containsKey(entityId)) {
      heartbeats.remove(entityId).cancel();
    }

    DynamoHeartbeat heartbeat = new DynamoHeartbeat(dynamo, lockTableName,
        heartbeatIntervalMs(), heartbeatTimeoutMs(), entityId, ownerId);
    heartbeat.schedule(scheduler());
    heartbeats.put(entityId, heartbeat);
  }

  @Override
  public boolean release(String entityId, String ownerId) {
    boolean succeeded = false;
    DynamoHeartbeat heartbeat = heartbeats.get(entityId);
    try {
      Tasks.foreach(entityId)
          .retry(RELEASE_RETRY_ATTEMPTS_MAX)
          .throwFailureWhenFinished()
          .onlyRetryOn(
              ProvisionedThroughputExceededException.class,
              TransactionConflictException.class,
              RequestLimitExceededException.class,
              InternalServerErrorException.class)
          .run(id -> dynamo.deleteItem(DeleteItemRequest.builder()
              .tableName(lockTableName)
              .key(toKey(id))
              .conditionExpression(CONDITION_LOCK_ID_MATCH)
              .expressionAttributeValues(toLockIdValues(id, ownerId))
              .build()));
      succeeded = true;
    } catch (ConditionalCheckFailedException e) {
      LOG.error("Failed to release lock for entity: {}, owner: {}, lock entity does not exist or owner not match",
          entityId, ownerId, e);
    } catch (DynamoDbException e) {
      LOG.error("Failed to release lock {} by for entity: {}, owner: {}, encountered unexpected DynamoDB exception",
          entityId, ownerId, e);
    } finally {
      if (heartbeat != null && heartbeat.ownerId().equals(ownerId)) {
        heartbeat.cancel();
      }
    }

    return succeeded;
  }

  private static Map toKey(String entityId) {
    return ImmutableMap.of(COL_LOCK_ENTITY_ID, AttributeValue.builder().s(entityId).build());
  }

  private static Map toNewItem(String entityId, String ownerId, long heartbeatTimeoutMs) {
    return ImmutableMap.of(
        COL_LOCK_ENTITY_ID, AttributeValue.builder().s(entityId).build(),
        COL_LOCK_OWNER_ID, AttributeValue.builder().s(ownerId).build(),
        COL_VERSION, AttributeValue.builder().s(UUID.randomUUID().toString()).build(),
        COL_LEASE_DURATION_MS, AttributeValue.builder().n(Long.toString(heartbeatTimeoutMs)).build());
  }

  private static Map toLockIdValues(String entityId, String ownerId) {
    return ImmutableMap.of(
        ":eid", AttributeValue.builder().s(entityId).build(),
        ":oid", AttributeValue.builder().s(ownerId).build());
  }

  @Override
  public void close() {
    dynamo.close();
    heartbeats.values().forEach(DynamoHeartbeat::cancel);
    heartbeats.clear();
  }

  /**
   * The lock table schema, for users who would like to create the table separately
   * @return lock table schema
   */
  public static List lockTableSchema() {
    return LOCK_TABLE_SCHEMA;
  }

  /**
   * The lock table column definition, for users who whould like to create the table separately
   * @return lock table column definition
   */
  public static List lockTableColDefinitions() {
    return LOCK_TABLE_COL_DEFINITIONS;
  }

  private static class DynamoHeartbeat implements Runnable {

    private final DynamoDbClient dynamo;
    private final String lockTableName;
    private final long intervalMs;
    private final long timeoutMs;
    private final String entityId;
    private final String ownerId;
    private ScheduledFuture future;

    DynamoHeartbeat(DynamoDbClient dynamo, String lockTableName, long intervalMs, long timeoutMs,
                           String entityId, String ownerId) {
      this.dynamo = dynamo;
      this.lockTableName = lockTableName;
      this.intervalMs = intervalMs;
      this.timeoutMs = timeoutMs;
      this.entityId = entityId;
      this.ownerId = ownerId;
      this.future = null;
    }

    @Override
    public void run() {
      try {
        dynamo.putItem(PutItemRequest.builder()
            .tableName(lockTableName)
            .item(toNewItem(entityId, ownerId, timeoutMs))
            .conditionExpression(CONDITION_LOCK_ID_MATCH)
            .expressionAttributeValues(toLockIdValues(entityId, ownerId))
            .build());
      } catch (ConditionalCheckFailedException e) {
        LOG.error("Fail to heartbeat for entity: {}, owner: {} due to conditional check failure, " +
                "unsafe concurrent commits might be going on", entityId, ownerId, e);
      } catch (RuntimeException e) {
        LOG.error("Failed to heartbeat for entity: {}, owner: {}", entityId, ownerId, e);
      }
    }

    public String ownerId() {
      return ownerId;
    }

    public void schedule(ScheduledExecutorService scheduler) {
      future = scheduler.scheduleAtFixedRate(this, 0, intervalMs, TimeUnit.MILLISECONDS);
    }

    public void cancel() {
      if (future != null) {
        future.cancel(false);
      }
    }
  }
}