You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hudi.apache.org by si...@apache.org on 2026/03/17 19:56:26 UTC
(hudi) branch release-0.14.2-prep updated: [HUDI-9159] Add support for StorageBasedLockProvider (#17970)
This is an automated email from the ASF dual-hosted git repository.
sivabalan pushed a commit to branch release-0.14.2-prep
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/release-0.14.2-prep by this push:
new a7a2cfde2955 [HUDI-9159] Add support for StorageBasedLockProvider (#17970)
a7a2cfde2955 is described below
commit a7a2cfde2955abcb03607458966d92d0fac71cfd
Author: Lokesh Jain <lj...@apache.org>
AuthorDate: Wed Mar 18 01:26:20 2026 +0530
[HUDI-9159] Add support for StorageBasedLockProvider (#17970)
[HUDI-9159] Add support for StorageBasedLockProvider
Cherry-picks the following PRs into release-0.14.2:
HUDI-9156: Add heartbeat manager for StorageBasedLockProvider
HUDI-9158: Add storage based lock provider abstract implementation and constructor support
HUDI-9160: Add GCS implementation of StorageBasedLockProvider
HUDI-9159: S3 implementation of StorageLock for StorageBasedLockProvider
HUDI-9122: Add RFC updates and tweak to S3 storage lock client
HUDI-9158: Rename config from hoodie.write.lock.storage.heartbeat.poll.secs to hoodie.write.lock.storage.renew.interval.secs
HUDI-9742: Add lock release metrics
HUDI-9752: Add metrics for storage based lock provider
HUDI-9782: Add audit logging scaffolding and implementation for storage based lock provider
HUDI-9168: Refactor hudi-client-common to not import hudi-aws
---------
Co-authored-by: Alex R <al...@onehouse.ai>
Co-authored-by: Alex R <ar...@gmail.com>
Co-authored-by: Y Ethan Guo <et...@gmail.com>
Co-authored-by: sivabalan <n....@gmail.com>
Co-authored-by: Lokesh Jain <lj...@192.168.1.11>
Co-authored-by: Jon Vexler <jb...@gmail.com>
Co-authored-by: Lokesh Jain <lj...@Lokeshs-MacBook-Pro.local>
Co-authored-by: Lokesh Jain <lj...@Lokeshs-MBP.lan>
Co-authored-by: Lokesh Jain <lj...@192.168.0.234>
Co-authored-by: Lokesh Jain <lj...@192.168.1.7>
Co-authored-by: Danny Chan <yu...@gmail.com>
Co-authored-by: Lokesh Jain <lj...@192.168.1.57>
Co-authored-by: Lin Liu <li...@gmail.com>
Co-authored-by: Lokesh Jain <lj...@192.168.1.8>
---
hudi-aws/pom.xml | 29 +
.../cloudwatch/CloudWatchMetricsReporter.java | 3 +-
.../cloudwatch/CloudWatchReporter.java | 5 +-
.../hudi/aws/sync/AWSGlueCatalogSyncClient.java | 2 +-
...amoDBBasedImplicitPartitionKeyLockProvider.java | 2 +-
.../aws/transaction/lock/S3StorageLockClient.java | 330 +++++++
.../cloudwatch/TestCloudWatchMetricsReporter.java | 21 +-
.../cloudwatch/TestCloudWatchReporter.java | 10 +-
.../lock/TestS3StorageBasedLockProvider.java | 160 ++++
.../transaction/lock/TestS3StorageLockClient.java | 336 +++++++
.../TestS3StorageLockClientAuditOperations.java | 272 ++++++
.../hudi/cli/commands/LockAuditingCommand.java | 661 ++++++++++++++
.../hudi/cli/commands/TestLockAuditingCommand.java | 945 ++++++++++++++++++++
hudi-client/hudi-client-common/pom.xml | 12 +-
.../lock/FileSystemBasedLockProvider.java | 5 +
.../hudi/client/transaction/lock/LockManager.java | 24 +-
.../transaction/lock/StorageBasedLockProvider.java | 651 ++++++++++++++
.../client/transaction/lock/StorageLockClient.java | 122 +++
...ZookeeperBasedImplicitBasePathLockProvider.java | 2 +-
.../lock/audit/AuditOperationState.java | 39 +
.../transaction/lock/audit/AuditService.java | 41 +
.../lock/audit/AuditServiceFactory.java | 119 +++
.../audit/StorageLockProviderAuditService.java | 177 ++++
.../lock/metrics/HoodieLockMetrics.java | 116 ++-
.../transaction/lock/models/HeartbeatManager.java | 49 ++
.../transaction/lock/models/LockGetResult.java | 38 +
.../lock/models/LockProviderHeartbeatManager.java | 351 ++++++++
.../transaction/lock/models/LockUpsertResult.java | 38 +
.../transaction/lock/models/StorageLockData.java | 71 ++
.../transaction/lock/models/StorageLockFile.java | 133 +++
.../apache/hudi/config/StorageBasedLockConfig.java | 101 +++
.../hudi/metrics/MetricsReporterFactory.java | 4 +-
.../client/transaction/TestHoodieLockMetrics.java | 348 ++++++++
.../hudi/client/transaction/TestLockManager.java | 47 +
.../lock/StorageBasedLockProviderTestBase.java | 290 +++++++
.../lock/TestStorageBasedLockProvider.java | 963 +++++++++++++++++++++
.../lock/audit/TestAuditServiceFactory.java | 160 ++++
.../audit/TestStorageLockProviderAuditService.java | 380 ++++++++
.../lock/models/StorageLockClientFileTest.java | 223 +++++
.../models/TestLockProviderHeartbeatManager.java | 481 ++++++++++
.../hudi/config/TestStorageBasedLockConfig.java | 101 +++
.../hudi/metrics/TestMetricsReporterFactory.java | 1 -
.../org/apache/hudi/common/fs/StorageSchemes.java | 77 +-
.../hudi/common/table/HoodieTableMetaClient.java | 2 +
.../org/apache/hudi/common/util/HoodieTimer.java | 11 +
.../java/org/apache/hudi/common/util/Option.java | 4 +
.../org/apache/hudi/common/fs/TestFSUtils.java | 20 +
.../apache/hudi/common/fs/TestStorageSchemes.java | 9 +
.../apache/hudi/common/util/TestHoodieTimer.java | 22 +-
hudi-gcp/pom.xml | 17 +
.../gcp/transaction/lock/GCSStorageLockClient.java | 289 +++++++
.../transaction/lock/TestGCSStorageLockClient.java | 304 +++++++
.../TestGCSStorageLockClientAuditOperations.java | 273 ++++++
.../hudi/command/procedures/BaseProcedure.scala | 7 +
.../procedures/CleanupAuditLockProcedure.scala | 190 ++++
.../hudi/command/procedures/HoodieProcedures.scala | 4 +
.../command/procedures/SetAuditLockProcedure.scala | 189 ++++
.../procedures/ShowAuditLockStatusProcedure.scala | 170 ++++
.../procedures/ValidateAuditLockProcedure.scala | 315 +++++++
.../procedure/TestCleanupAuditLockProcedure.scala | 294 +++++++
.../hudi/procedure/TestSetAuditLockProcedure.scala | 169 ++++
.../TestShowAuditLockStatusProcedure.scala | 223 +++++
.../procedure/TestValidateAuditLockProcedure.scala | 454 ++++++++++
pom.xml | 17 +-
64 files changed, 10850 insertions(+), 73 deletions(-)
diff --git a/hudi-aws/pom.xml b/hudi-aws/pom.xml
index 3093134d3a09..bc280da053d3 100644
--- a/hudi-aws/pom.xml
+++ b/hudi-aws/pom.xml
@@ -41,6 +41,11 @@
</dependency>
<!-- Hoodie -->
+ <dependency>
+ <groupId>org.apache.hudi</groupId>
+ <artifactId>hudi-client-common</artifactId>
+ <version>${project.version}</version>
+ </dependency>
<dependency>
<groupId>org.apache.hudi</groupId>
<artifactId>hudi-common</artifactId>
@@ -168,6 +173,13 @@
<version>${aws.sdk.version}</version>
</dependency>
+ <!-- https://mvnrepository.com/artifact/com.amazonaws/aws-java-sdk-s3 -->
+ <dependency>
+ <groupId>software.amazon.awssdk</groupId>
+ <artifactId>s3</artifactId>
+ <version>${aws.sdk.version}</version>
+ </dependency>
+
<!-- https://mvnrepository.com/artifact/com.amazonaws/aws-java-sdk-sqs -->
<dependency>
<groupId>software.amazon.awssdk</groupId>
@@ -193,6 +205,23 @@
<version>${project.version}</version>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.hudi</groupId>
+ <artifactId>hudi-client-common</artifactId>
+ <version>${project.version}</version>
+ <classifier>tests</classifier>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.testcontainers</groupId>
+ <artifactId>testcontainers</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.testcontainers</groupId>
+ <artifactId>localstack</artifactId>
+ <scope>test</scope>
+ </dependency>
</dependencies>
<build>
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/cloudwatch/CloudWatchMetricsReporter.java b/hudi-aws/src/main/java/org/apache/hudi/aws/metrics/cloudwatch/CloudWatchMetricsReporter.java
similarity index 96%
rename from hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/cloudwatch/CloudWatchMetricsReporter.java
rename to hudi-aws/src/main/java/org/apache/hudi/aws/metrics/cloudwatch/CloudWatchMetricsReporter.java
index d05632b9bbf8..f3fd9d60b5c6 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/cloudwatch/CloudWatchMetricsReporter.java
+++ b/hudi-aws/src/main/java/org/apache/hudi/aws/metrics/cloudwatch/CloudWatchMetricsReporter.java
@@ -16,9 +16,8 @@
* limitations under the License.
*/
-package org.apache.hudi.metrics.cloudwatch;
+package org.apache.hudi.aws.metrics.cloudwatch;
-import org.apache.hudi.aws.cloudwatch.CloudWatchReporter;
import org.apache.hudi.config.HoodieWriteConfig;
import org.apache.hudi.metrics.MetricsReporter;
diff --git a/hudi-aws/src/main/java/org/apache/hudi/aws/cloudwatch/CloudWatchReporter.java b/hudi-aws/src/main/java/org/apache/hudi/aws/metrics/cloudwatch/CloudWatchReporter.java
similarity index 98%
rename from hudi-aws/src/main/java/org/apache/hudi/aws/cloudwatch/CloudWatchReporter.java
rename to hudi-aws/src/main/java/org/apache/hudi/aws/metrics/cloudwatch/CloudWatchReporter.java
index f8fefbc69348..c440159b4fe5 100644
--- a/hudi-aws/src/main/java/org/apache/hudi/aws/cloudwatch/CloudWatchReporter.java
+++ b/hudi-aws/src/main/java/org/apache/hudi/aws/metrics/cloudwatch/CloudWatchReporter.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.hudi.aws.cloudwatch;
+package org.apache.hudi.aws.metrics.cloudwatch;
import org.apache.hudi.aws.credentials.HoodieAWSCredentialsProviderFactory;
import org.apache.hudi.common.util.Option;
@@ -235,6 +235,9 @@ public class CloudWatchReporter extends ScheduledReporter {
} catch (final Exception ex) {
LOG.error("Error reporting metrics to CloudWatch. The data in this CloudWatch request "
+ "may have been discarded, and not made it to CloudWatch.", ex);
+ if (ex instanceof InterruptedException) {
+ Thread.currentThread().interrupt();
+ }
}
}
}
diff --git a/hudi-aws/src/main/java/org/apache/hudi/aws/sync/AWSGlueCatalogSyncClient.java b/hudi-aws/src/main/java/org/apache/hudi/aws/sync/AWSGlueCatalogSyncClient.java
index 049d633c0257..e480df00ad29 100644
--- a/hudi-aws/src/main/java/org/apache/hudi/aws/sync/AWSGlueCatalogSyncClient.java
+++ b/hudi-aws/src/main/java/org/apache/hudi/aws/sync/AWSGlueCatalogSyncClient.java
@@ -71,7 +71,7 @@ import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;
-import static org.apache.hudi.aws.utils.S3Utils.s3aToS3;
+import static org.apache.hudi.common.fs.FSUtils.s3aToS3;
import static org.apache.hudi.common.util.MapUtils.containsAll;
import static org.apache.hudi.common.util.MapUtils.isNullOrEmpty;
import static org.apache.hudi.config.GlueCatalogSyncClientConfig.GLUE_METADATA_FILE_LISTING;
diff --git a/hudi-aws/src/main/java/org/apache/hudi/aws/transaction/lock/DynamoDBBasedImplicitPartitionKeyLockProvider.java b/hudi-aws/src/main/java/org/apache/hudi/aws/transaction/lock/DynamoDBBasedImplicitPartitionKeyLockProvider.java
index c9ea367493d7..1b34f7974bab 100644
--- a/hudi-aws/src/main/java/org/apache/hudi/aws/transaction/lock/DynamoDBBasedImplicitPartitionKeyLockProvider.java
+++ b/hudi-aws/src/main/java/org/apache/hudi/aws/transaction/lock/DynamoDBBasedImplicitPartitionKeyLockProvider.java
@@ -31,7 +31,7 @@ import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import javax.annotation.concurrent.NotThreadSafe;
-import static org.apache.hudi.aws.utils.S3Utils.s3aToS3;
+import static org.apache.hudi.common.fs.FSUtils.s3aToS3;
/**
* A DynamoDB based lock.
diff --git a/hudi-aws/src/main/java/org/apache/hudi/aws/transaction/lock/S3StorageLockClient.java b/hudi-aws/src/main/java/org/apache/hudi/aws/transaction/lock/S3StorageLockClient.java
new file mode 100644
index 000000000000..67fef4e75785
--- /dev/null
+++ b/hudi-aws/src/main/java/org/apache/hudi/aws/transaction/lock/S3StorageLockClient.java
@@ -0,0 +1,330 @@
+/*
+ * 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.hudi.aws.transaction.lock;
+
+import org.apache.hudi.aws.credentials.HoodieAWSCredentialsProviderFactory;
+import org.apache.hudi.client.transaction.lock.StorageLockClient;
+import org.apache.hudi.client.transaction.lock.models.LockGetResult;
+import org.apache.hudi.client.transaction.lock.models.LockUpsertResult;
+import org.apache.hudi.client.transaction.lock.models.StorageLockData;
+import org.apache.hudi.client.transaction.lock.models.StorageLockFile;
+import org.apache.hudi.common.util.Functions;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.common.util.VisibleForTesting;
+import org.apache.hudi.common.util.collection.Pair;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.core.ResponseInputStream;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.GetBucketLocationRequest;
+import software.amazon.awssdk.services.s3.model.GetBucketLocationResponse;
+import software.amazon.awssdk.services.s3.model.GetObjectRequest;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.services.s3.model.S3Exception;
+
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.time.Duration;
+import java.util.Properties;
+
+import static org.apache.hudi.config.StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS;
+
+/**
+ * S3-based distributed lock client using ETag checks (AWS SDK v2).
+ * See RFC: <a href="https://github.com/apache/hudi/blob/master/rfc/rfc-91/rfc-91.md">RFC-91</a>
+ */
+@ThreadSafe
+public class S3StorageLockClient implements StorageLockClient {
+
+ private static final Logger LOG = LoggerFactory.getLogger(S3StorageLockClient.class);
+ private static final int PRECONDITION_FAILURE_ERROR_CODE = 412;
+ private static final int NOT_FOUND_ERROR_CODE = 404;
+ private static final int CONDITIONAL_REQUEST_CONFLICT_ERROR_CODE = 409;
+ private static final int RATE_LIMIT_ERROR_CODE = 429;
+ private static final int INTERNAL_SERVER_ERROR_CODE_MIN = 500;
+
+ private final Logger logger;
+ private final S3Client s3Client;
+ private final String bucketName;
+ private final String lockFilePath;
+ private final String ownerId;
+
+ /**
+ * Constructor that is used by reflection to instantiate an S3-based storage locking client.
+ *
+ * @param ownerId The owner id.
+ * @param lockFileUri The full table base path where the lock will be written.
+ * @param props The properties for the lock config, can be used to customize client.
+ */
+ public S3StorageLockClient(String ownerId, String lockFileUri, Properties props) {
+ this(ownerId, lockFileUri, props, createDefaultS3Client(), LOG);
+ }
+
+ @VisibleForTesting
+ S3StorageLockClient(String ownerId, String lockFileUri, Properties props, Functions.Function2<String, Properties, S3Client> s3ClientSupplier, Logger logger) {
+ Pair<String, String> bucketAndPath = StorageLockClient.parseBucketAndPath(lockFileUri);
+ this.bucketName = bucketAndPath.getLeft();
+ this.lockFilePath = bucketAndPath.getRight();
+
+ this.s3Client = s3ClientSupplier.apply(bucketName, props);
+ this.ownerId = ownerId;
+ this.logger = logger;
+ }
+
+ @Override
+ public Pair<LockGetResult, Option<StorageLockFile>> readCurrentLockFile() {
+ try (ResponseInputStream<GetObjectResponse> in = s3Client.getObject(
+ GetObjectRequest.builder()
+ .bucket(bucketName)
+ .key(lockFilePath)
+ .build())) {
+ String eTag = in.response().eTag();
+ return Pair.of(LockGetResult.SUCCESS, Option.of(StorageLockFile.createFromStream(in, eTag)));
+ } catch (S3Exception e) {
+ int status = e.statusCode();
+ // Default to unknown error
+ LockGetResult result = LockGetResult.UNKNOWN_ERROR;
+ if (status == NOT_FOUND_ERROR_CODE) {
+ logger.info("OwnerId: {}, Object not found: {}", ownerId, lockFilePath);
+ result = LockGetResult.NOT_EXISTS;
+ } else if (status == CONDITIONAL_REQUEST_CONFLICT_ERROR_CODE) {
+ logger.info("OwnerId: {}, Conflicting operation has occurred: {}", ownerId, lockFilePath);
+ } else if (status == RATE_LIMIT_ERROR_CODE) {
+ logger.warn("OwnerId: {}, Rate limit exceeded: {}", ownerId, lockFilePath);
+ } else if (status >= INTERNAL_SERVER_ERROR_CODE_MIN) {
+ logger.warn("OwnerId: {}, S3 internal server error: {}", ownerId, lockFilePath, e);
+ } else {
+ throw e;
+ }
+ return Pair.of(result, Option.empty());
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed reading lock file from S3: " + lockFilePath, e);
+ }
+ }
+
+ @Override
+ public Pair<LockUpsertResult, Option<StorageLockFile>> tryUpsertLockFile(StorageLockData newLockData, Option<StorageLockFile> previousLockFile) {
+ boolean isLockRenewal = previousLockFile.isPresent();
+ String currentEtag = isLockRenewal ? previousLockFile.get().getVersionId() : null;
+ LockUpsertResult result = LockUpsertResult.UNKNOWN_ERROR;
+ try {
+ StorageLockFile updated = createOrUpdateLockFileInternal(newLockData, currentEtag);
+ return Pair.of(LockUpsertResult.SUCCESS, Option.of(updated));
+ } catch (S3Exception e) {
+ result = handleUpsertS3Exception(e);
+ } catch (AwsServiceException | SdkClientException e) {
+ logger.warn("OwnerId: {}, Unexpected SDK error while writing lock file: {}", ownerId, lockFilePath, e);
+ if (!isLockRenewal) {
+ // We should always throw errors early when we are creating the lock file.
+ // This is likely indicative of a larger issue that should bubble up sooner.
+ throw e;
+ }
+ }
+
+ return Pair.of(result, Option.empty());
+ }
+
+ /**
+ * Internal helper to create or update the lock file with optional ETag precondition.
+ */
+ private StorageLockFile createOrUpdateLockFileInternal(StorageLockData lockData, String expectedEtag) {
+ byte[] bytes = StorageLockFile.toByteArray(lockData);
+ PutObjectRequest.Builder putRequestBuilder = PutObjectRequest.builder()
+ .bucket(bucketName)
+ .key(lockFilePath);
+
+ // ETag-based constraints:
+ // - If expectedEtag is not null:
+ // We assume that the file already exists on S3 with the ETag "expectedEtag".
+ // The update operation will include an ifMatch(expectedEtag) condition, meaning the update will only
+ // succeed if the current file's ETag exactly matches expectedEtag.
+ // If the actual ETag of the file on S3 differs from expectedEtag, the update attempt will fail.
+ // - If expectedEtag is null:
+ // We assume that the file does not currently exist on S3.
+ // The operation will use ifNoneMatch("*"), which instructs S3 to create the file only if it doesn't already exist.
+ // If a file with the same name is present (i.e., there is an existing ETag), the creation attempt will fail.
+ if (expectedEtag == null) {
+ putRequestBuilder.ifNoneMatch("*");
+ } else {
+ putRequestBuilder.ifMatch(expectedEtag);
+ }
+
+ PutObjectResponse response = s3Client.putObject(putRequestBuilder.build(), RequestBody.fromBytes(bytes));
+ String newEtag = response.eTag();
+
+ return new StorageLockFile(lockData, newEtag);
+ }
+
+ private LockUpsertResult handleUpsertS3Exception(S3Exception e) {
+ int status = e.statusCode();
+ if (status == PRECONDITION_FAILURE_ERROR_CODE) {
+ logger.warn("OwnerId: {}, Lockfile modified by another process: {}", ownerId, lockFilePath);
+ return LockUpsertResult.ACQUIRED_BY_OTHERS;
+ } else if (status == CONDITIONAL_REQUEST_CONFLICT_ERROR_CODE) {
+ logger.warn("OwnerId: {}, Retriable conditional request conflict error: {}", ownerId, lockFilePath);
+ } else if (status == RATE_LIMIT_ERROR_CODE) {
+ logger.warn("OwnerId: {}, Rate limit exceeded for: {}", ownerId, lockFilePath);
+ } else if (status >= INTERNAL_SERVER_ERROR_CODE_MIN) {
+ logger.warn("OwnerId: {}, internal server error for: {}", ownerId, lockFilePath, e);
+ } else {
+ logger.error("OwnerId: {}, Error writing lock file: {}", ownerId, lockFilePath, e);
+ }
+
+ return LockUpsertResult.UNKNOWN_ERROR;
+ }
+
+ private static Functions.Function2<String, Properties, S3Client> createDefaultS3Client() {
+ return (bucketName, props) -> {
+ Region region;
+ boolean requiredFallbackRegion = false;
+ try {
+ region = DefaultAwsRegionProviderChain.builder().build().getRegion();
+ } catch (SdkClientException e) {
+ // Fallback to us-east-1 if no region is found
+ region = Region.US_EAST_1;
+ requiredFallbackRegion = true;
+ }
+
+ // Set all request timeouts to be 1/5 of the default validity.
+ // Each call to acquire a lock requires 2 requests.
+ // Each renewal requires 1 request.
+ long validityTimeoutSecs =
+ ((Number) props.getOrDefault(
+ VALIDITY_TIMEOUT_SECONDS.key(),
+ VALIDITY_TIMEOUT_SECONDS.defaultValue()
+ )).longValue();
+ long s3CallTimeoutSecs = validityTimeoutSecs / 5;
+ S3Client s3Client = createS3Client(region, s3CallTimeoutSecs, props);
+ if (requiredFallbackRegion) {
+ GetBucketLocationResponse bucketLocationResponse = s3Client.getBucketLocation(
+ GetBucketLocationRequest.builder().bucket(bucketName).build());
+ // This is null when the region is US_EAST_1, so we do not need to worry about duplicate logic.
+ String regionString = bucketLocationResponse.locationConstraintAsString();
+ if (!StringUtils.isNullOrEmpty(regionString)) {
+ // Close existing client and create another.
+ s3Client.close();
+ return createS3Client(
+ Region.of(regionString),
+ s3CallTimeoutSecs,
+ props);
+ }
+ }
+
+ return s3Client;
+ };
+ }
+
+ private static S3Client createS3Client(Region region, long timeoutSecs, Properties props) {
+ // Set the timeout, credentials, and region
+ return S3Client.builder()
+ .overrideConfiguration(
+ b -> b.apiCallTimeout(Duration.ofSeconds(timeoutSecs)))
+ .credentialsProvider(HoodieAWSCredentialsProviderFactory.getAwsCredentialsProvider(props))
+ .region(region).build();
+ }
+
+ @Override
+ public Option<String> readObject(String filePath, boolean checkExistsFirst) {
+ try {
+ // Parse the file path to get bucket and key
+ Pair<String, String> bucketAndKey = StorageLockClient.parseBucketAndPath(filePath);
+ String bucket = bucketAndKey.getLeft();
+ String key = bucketAndKey.getRight();
+
+ if (checkExistsFirst) {
+ // First check if the file exists (lightweight HEAD request)
+ try {
+ s3Client.headObject(HeadObjectRequest.builder()
+ .bucket(bucket)
+ .key(key)
+ .build());
+ } catch (S3Exception e) {
+ if (e.statusCode() == NOT_FOUND_ERROR_CODE) {
+ // File doesn't exist - this is the common case for optional configs
+ logger.debug("JSON config file not found: {}", filePath);
+ return Option.empty();
+ }
+ throw e; // Re-throw other errors
+ }
+ }
+
+ // Read the file (either after existence check or directly)
+ String content = s3Client.getObjectAsBytes(
+ GetObjectRequest.builder()
+ .bucket(bucket)
+ .key(key)
+ .build()).asUtf8String();
+
+ return Option.of(content);
+ } catch (S3Exception e) {
+ if (e.statusCode() == NOT_FOUND_ERROR_CODE) {
+ logger.debug("JSON config file not found: {}", filePath);
+ return Option.empty();
+ }
+ logger.warn("Error reading JSON config file: {}", filePath, e);
+ return Option.empty();
+ } catch (Exception e) {
+ logger.warn("Error reading JSON config file: {}", filePath, e);
+ return Option.empty();
+ }
+ }
+
+ @Override
+ public boolean writeObject(String filePath, String content) {
+ try {
+ // Parse the file path to get bucket and key
+ Pair<String, String> bucketAndPath = StorageLockClient.parseBucketAndPath(filePath);
+ String bucket = bucketAndPath.getLeft();
+ String key = bucketAndPath.getRight();
+
+ // Write the content to S3
+ s3Client.putObject(
+ PutObjectRequest.builder()
+ .bucket(bucket)
+ .key(key)
+ .build(),
+ RequestBody.fromString(content));
+
+ logger.debug("Successfully wrote object to: {}", filePath);
+ return true;
+ } catch (Exception e) {
+ logger.warn("Error writing object to: {}", filePath, e);
+ return false;
+ }
+ }
+
+ @Override
+ public void close() {
+ s3Client.close();
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/cloudwatch/TestCloudWatchMetricsReporter.java b/hudi-aws/src/test/java/org/apache/hudi/aws/metrics/cloudwatch/TestCloudWatchMetricsReporter.java
similarity index 76%
rename from hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/cloudwatch/TestCloudWatchMetricsReporter.java
rename to hudi-aws/src/test/java/org/apache/hudi/aws/metrics/cloudwatch/TestCloudWatchMetricsReporter.java
index 7901d8024651..52ebb5e1d2db 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/cloudwatch/TestCloudWatchMetricsReporter.java
+++ b/hudi-aws/src/test/java/org/apache/hudi/aws/metrics/cloudwatch/TestCloudWatchMetricsReporter.java
@@ -16,9 +16,8 @@
* limitations under the License.
*/
-package org.apache.hudi.metrics.cloudwatch;
+package org.apache.hudi.aws.metrics.cloudwatch;
-import org.apache.hudi.aws.cloudwatch.CloudWatchReporter;
import org.apache.hudi.config.HoodieWriteConfig;
import com.codahale.metrics.MetricRegistry;
@@ -36,6 +35,9 @@ import static org.mockito.Mockito.when;
@ExtendWith(MockitoExtension.class)
public class TestCloudWatchMetricsReporter {
+ @Mock
+ private HoodieWriteConfig writeConfig;
+
@Mock
private HoodieWriteConfig config;
@@ -59,4 +61,19 @@ public class TestCloudWatchMetricsReporter {
metricsReporter.stop();
verify(reporter, times(1)).stop();
}
+
+ @Test
+ public void testReporterUsingMetricsConfig() {
+ when(writeConfig.getCloudWatchReportPeriodSeconds()).thenReturn(30);
+ CloudWatchMetricsReporter metricsReporter = new CloudWatchMetricsReporter(writeConfig, registry, reporter);
+
+ metricsReporter.start();
+ verify(reporter, times(1)).start(30, TimeUnit.SECONDS);
+
+ metricsReporter.report();
+ verify(reporter, times(1)).report();
+
+ metricsReporter.stop();
+ verify(reporter, times(1)).stop();
+ }
}
diff --git a/hudi-aws/src/test/java/org/apache/hudi/aws/cloudwatch/TestCloudWatchReporter.java b/hudi-aws/src/test/java/org/apache/hudi/aws/metrics/cloudwatch/TestCloudWatchReporter.java
similarity index 94%
rename from hudi-aws/src/test/java/org/apache/hudi/aws/cloudwatch/TestCloudWatchReporter.java
rename to hudi-aws/src/test/java/org/apache/hudi/aws/metrics/cloudwatch/TestCloudWatchReporter.java
index 600dd5786959..dc3a4f0493c4 100644
--- a/hudi-aws/src/test/java/org/apache/hudi/aws/cloudwatch/TestCloudWatchReporter.java
+++ b/hudi-aws/src/test/java/org/apache/hudi/aws/metrics/cloudwatch/TestCloudWatchReporter.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.hudi.aws.cloudwatch;
+package org.apache.hudi.aws.metrics.cloudwatch;
import software.amazon.awssdk.services.cloudwatch.CloudWatchAsyncClient;
import software.amazon.awssdk.services.cloudwatch.model.Dimension;
@@ -49,10 +49,10 @@ import java.util.TreeMap;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
-import static org.apache.hudi.aws.cloudwatch.CloudWatchReporter.DIMENSION_COUNT_TYPE_VALUE;
-import static org.apache.hudi.aws.cloudwatch.CloudWatchReporter.DIMENSION_GAUGE_TYPE_VALUE;
-import static org.apache.hudi.aws.cloudwatch.CloudWatchReporter.DIMENSION_METRIC_TYPE_KEY;
-import static org.apache.hudi.aws.cloudwatch.CloudWatchReporter.DIMENSION_TABLE_NAME_KEY;
+import static org.apache.hudi.aws.metrics.cloudwatch.CloudWatchReporter.DIMENSION_COUNT_TYPE_VALUE;
+import static org.apache.hudi.aws.metrics.cloudwatch.CloudWatchReporter.DIMENSION_GAUGE_TYPE_VALUE;
+import static org.apache.hudi.aws.metrics.cloudwatch.CloudWatchReporter.DIMENSION_METRIC_TYPE_KEY;
+import static org.apache.hudi.aws.metrics.cloudwatch.CloudWatchReporter.DIMENSION_TABLE_NAME_KEY;
import static org.junit.jupiter.api.Assertions.assertEquals;
@ExtendWith(MockitoExtension.class)
diff --git a/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageBasedLockProvider.java b/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageBasedLockProvider.java
new file mode 100644
index 000000000000..343bf26bc1fe
--- /dev/null
+++ b/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageBasedLockProvider.java
@@ -0,0 +1,160 @@
+/*
+ * 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.hudi.aws.transaction.lock;
+
+import org.apache.hudi.client.transaction.lock.StorageBasedLockProviderTestBase;
+import org.apache.hudi.client.transaction.lock.StorageBasedLockProvider;
+import org.apache.hudi.common.config.LockConfiguration;
+import org.apache.hudi.config.HoodieAWSConfig;
+
+import org.apache.hadoop.conf.Configuration;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.testcontainers.containers.localstack.LocalStackContainer;
+import org.testcontainers.utility.DockerImageName;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.S3ClientBuilder;
+import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.services.s3.model.ServerSideEncryption;
+
+import static org.apache.hudi.config.HoodieWriteConfig.BASE_PATH;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.mockito.Mockito.mockStatic;
+import static org.testcontainers.containers.localstack.LocalStackContainer.Service.S3;
+
+/**
+ * Tests S3-based StorageBasedLockProvider using a LocalStack container
+ * to emulate S3.
+ */
+@Disabled("HUDI-9159 The tests do not work. Disabling them to unblock Azure CI")
+public class TestS3StorageBasedLockProvider extends StorageBasedLockProviderTestBase {
+
+ private static final DockerImageName LOCALSTACK_IMAGE =
+ DockerImageName.parse("localstack/localstack:4.1.0");
+
+ private static LocalStackContainer S3_CONTAINER;
+
+ private static final String TEST_BUCKET = "test-bucket";
+ private static S3Client s3Client;
+
+ @BeforeAll
+ static void initContainer() {
+ // Spin up LocalStack with S3 service
+ S3_CONTAINER = new LocalStackContainer(LOCALSTACK_IMAGE)
+ .withServices(S3);
+ S3_CONTAINER.start();
+
+ s3Client = S3Client.builder()
+ .endpointOverride(S3_CONTAINER.getEndpointOverride(S3))
+ .credentialsProvider(
+ StaticCredentialsProvider.create(
+ AwsBasicCredentials.create(
+ S3_CONTAINER.getAccessKey(),
+ S3_CONTAINER.getSecretKey()
+ )
+ )
+ )
+ .region(Region.of(S3_CONTAINER.getRegion()))
+ .build();
+ s3Client.createBucket(CreateBucketRequest.builder()
+ .bucket(TEST_BUCKET)
+ .build());
+ }
+
+ @AfterAll
+ static void stopContainer() {
+ if (S3_CONTAINER != null) {
+ S3_CONTAINER.stop();
+ }
+ }
+
+ /**
+ * Creates a real S3-based StorageBasedLockProvider for testing.
+ */
+ @Override
+ protected StorageBasedLockProvider createLockProvider() {
+ LockConfiguration lockConf = new LockConfiguration(providerProperties);
+
+ try (MockedStatic<S3Client> s3ClientStaticMock =
+ mockStatic(S3Client.class, Mockito.CALLS_REAL_METHODS)) {
+ s3ClientStaticMock.when(S3Client::builder).thenAnswer(invocation -> {
+ S3ClientBuilder builder = (S3ClientBuilder) invocation.callRealMethod();
+ // Always add the endpoint override from S3_CONTAINER
+ builder.endpointOverride(S3_CONTAINER.getEndpointOverride(S3));
+ return builder;
+ });
+ return new StorageBasedLockProvider(
+ lockConf,
+ (Configuration) null);
+ }
+ }
+
+ @BeforeEach
+ void setupLockProvider() {
+ // Provide base path for the test
+ providerProperties.put(BASE_PATH.key(), String.format("s3://%s/lake/db/tbl-default", TEST_BUCKET));
+ providerProperties.put(HoodieAWSConfig.AWS_SECRET_KEY.key(), S3_CONTAINER.getSecretKey());
+ providerProperties.put(HoodieAWSConfig.AWS_ACCESS_KEY.key(), S3_CONTAINER.getAccessKey());
+ lockProvider = createLockProvider();
+ }
+
+ @Test
+ void testS3Preconditions() {
+ // Similar to GCS test, just to show we can put and head objects
+ String key = "my-test-object";
+ // Put an object
+ PutObjectResponse put1 = s3Client.putObject(
+ PutObjectRequest.builder()
+ .bucket(TEST_BUCKET)
+ .key(key)
+ .serverSideEncryption(ServerSideEncryption.AES256)
+ .ifNoneMatch("*")
+ .build(),
+ RequestBody.fromString("Hello1"));
+
+ // Put a second object with different contents
+ PutObjectResponse put2 = s3Client.putObject(
+ PutObjectRequest.builder()
+ .bucket(TEST_BUCKET)
+ .key(key)
+ .serverSideEncryption(ServerSideEncryption.AES256)
+ .ifMatch(put1.eTag())
+ .build(),
+ RequestBody.fromString("Hello2"));
+ s3Client.close();
+ // ETags should differ
+ assertNotNull(put1.eTag());
+ assertNotNull(put2.eTag());
+ assertNotEquals(put1.eTag(), put2.eTag(), "ETags should differ if the content changed");
+ }
+}
diff --git a/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageLockClient.java b/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageLockClient.java
new file mode 100644
index 000000000000..1c81c58327cb
--- /dev/null
+++ b/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageLockClient.java
@@ -0,0 +1,336 @@
+/*
+ * 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.hudi.aws.transaction.lock;
+
+import org.apache.hudi.client.transaction.lock.models.LockGetResult;
+import org.apache.hudi.client.transaction.lock.models.LockUpsertResult;
+import org.apache.hudi.client.transaction.lock.models.StorageLockData;
+import org.apache.hudi.client.transaction.lock.models.StorageLockFile;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.exception.HoodieIOException;
+
+import org.apache.hudi.exception.HoodieLockException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.slf4j.Logger;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.core.ResponseInputStream;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.GetObjectRequest;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.services.s3.model.S3Exception;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Properties;
+
+import static org.apache.hudi.client.transaction.lock.models.LockUpsertResult.UNKNOWN_ERROR;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.contains;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+class TestS3StorageLockClient {
+
+ private static final String OWNER_ID = "ownerId";
+ private static final String LOCK_FILE_URI = "s3://bucket/lockFilePath";
+ private static final String LOCK_FILE_PATH = "lockFilePath";
+
+ @Mock
+ private S3Client mockS3Client;
+
+ @Mock
+ private Logger mockLogger;
+
+ private S3StorageLockClient lockService;
+
+ @BeforeEach
+ void setUp() {
+ lockService = new S3StorageLockClient(
+ OWNER_ID,
+ LOCK_FILE_URI,
+ new Properties(),
+ (a,b) -> mockS3Client,
+ mockLogger
+ );
+ }
+
+ private void mockS3ObjectWithLockData(StorageLockData lockData, String eTag) {
+ byte[] bytes = StorageLockFile.toByteArray(lockData);
+ ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
+
+ ResponseInputStream<GetObjectResponse> mockResponseStream = new ResponseInputStream<>(
+ GetObjectResponse.builder().eTag(eTag).build(),
+ new MockAbortableInputStream(bais)
+ );
+
+ when(mockS3Client.getObject(any(GetObjectRequest.class))).thenReturn(mockResponseStream);
+ }
+
+ private static class MockAbortableInputStream extends java.io.FilterInputStream {
+ protected MockAbortableInputStream(ByteArrayInputStream in) {
+ super(in);
+ }
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_noPreviousLock_success() {
+ StorageLockData lockData = new StorageLockData(false, System.currentTimeMillis(), "myTxOwner");
+ PutObjectResponse putResp = PutObjectResponse.builder().eTag("new-etag-123").build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenReturn(putResp);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result = lockService.tryUpsertLockFile(lockData, Option.empty());
+
+ assertEquals(LockUpsertResult.SUCCESS, result.getLeft());
+ assertTrue(result.getRight().isPresent());
+ assertEquals("new-etag-123", result.getRight().get().getVersionId());
+ verify(mockS3Client, times(1)).putObject(any(PutObjectRequest.class), any(RequestBody.class));
+ verifyNoMoreInteractions(mockLogger);
+ }
+
+ @Test
+ void testInitializeWithInvalidUri() {
+ assertThrows(HoodieLockException.class, () -> new S3StorageLockClient(
+ OWNER_ID,
+ "\\",
+ new Properties(),
+ (a, b) -> mockS3Client,
+ mockLogger
+ ));
+ }
+
+ @Test
+ void testInitializeWithNoLockFilePath() {
+ IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> new S3StorageLockClient(
+ OWNER_ID,
+ "s3://bucket/",
+ new Properties(),
+ (a, b) -> mockS3Client,
+ mockLogger
+ ));
+ assertTrue(ex.getMessage().contains("path"));
+ }
+
+ @Test
+ void testInitializeWithNoBucketName() {
+ IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> new S3StorageLockClient(
+ OWNER_ID,
+ "s3:///path",
+ new Properties(),
+ (a, b) -> mockS3Client,
+ mockLogger
+ ));
+ assertTrue(ex.getMessage().contains("bucket name"));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_withPreviousLock_success() {
+ StorageLockData lockData = new StorageLockData(false, 1000L, "myTxOwner");
+ StorageLockFile prevLockFile = new StorageLockFile(lockData, "old-etag-999");
+ PutObjectResponse putResp = PutObjectResponse.builder().eTag("new-etag-456").build();
+ when(mockS3Client.putObject(
+ eq(PutObjectRequest.builder().bucket("bucket")
+ .key(LOCK_FILE_PATH)
+ .ifMatch("old-etag-999")
+ .build()),
+ any(RequestBody.class)
+ )).thenReturn(putResp);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result =
+ lockService.tryUpsertLockFile(lockData, Option.of(prevLockFile));
+
+ assertEquals(LockUpsertResult.SUCCESS, result.getLeft());
+ assertTrue(result.getRight().isPresent());
+ assertEquals("new-etag-456", result.getRight().get().getVersionId());
+ verify(mockS3Client, times(1)).putObject(any(PutObjectRequest.class), any(RequestBody.class));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_preconditionFailed() {
+ StorageLockData lockData = new StorageLockData(false, 2000L, "txOwner");
+ StorageLockFile prevLockFile = new StorageLockFile(lockData, "some-etag");
+
+ AwsServiceException ex412 = S3Exception.builder().statusCode(412).build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenThrow(ex412);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result =
+ lockService.tryUpsertLockFile(lockData, Option.of(prevLockFile));
+
+ assertEquals(LockUpsertResult.ACQUIRED_BY_OTHERS, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).warn(contains("Lockfile modified by another process"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_conflict409() {
+ StorageLockData lockData = new StorageLockData(false, 3000L, "myTxOwner");
+ StorageLockFile prevLockFile = new StorageLockFile(lockData, "some-etag-409");
+
+ AwsServiceException ex409 = S3Exception.builder().statusCode(409).build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenThrow(ex409);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result =
+ lockService.tryUpsertLockFile(lockData, Option.of(prevLockFile));
+
+ // This error is not necessarily unknown, but we don't know the state of the lockfile,
+ // which is why the return result is UNKNOWN_ERROR
+ assertEquals(UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).warn(contains("Retriable conditional request conflict error"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_rateLimitExceeded() {
+ StorageLockData lockData = new StorageLockData(false, 4000L, "myTxOwner");
+ AwsServiceException ex429 = S3Exception.builder().statusCode(429).build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenThrow(ex429);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result =
+ lockService.tryUpsertLockFile(lockData, Option.empty());
+
+ assertEquals(UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).warn(contains("Rate limit exceeded"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_serverError() {
+ StorageLockData lockData = new StorageLockData(false, 5000L, "myTxOwner");
+ AwsServiceException ex503 = S3Exception.builder().statusCode(503).build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenThrow(ex503);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result =
+ lockService.tryUpsertLockFile(lockData, Option.empty());
+
+ assertEquals(UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).warn(contains("internal server error"), eq(OWNER_ID), eq(LOCK_FILE_PATH), eq(ex503));
+ }
+
+ @Test
+ void testTryCreateLockFile_unexpectedError() {
+ StorageLockData lockData = new StorageLockData(false, 8000L, "myTxOwner");
+ AwsServiceException ex400 = AwsServiceException.builder().statusCode(400).build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenThrow(ex400);
+
+ assertThrows(AwsServiceException.class, () -> lockService.tryUpsertLockFile(lockData, Option.empty()));
+ }
+
+ @Test
+ void testGetCurrentLockFile_404Error() {
+ AwsServiceException ex404 = S3Exception.builder().statusCode(404).build();
+ when(mockS3Client.getObject(any(GetObjectRequest.class))).thenThrow(ex404);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+ assertEquals(LockGetResult.NOT_EXISTS, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).info(contains("Object not found"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testGetCurrentLockFile_409Error() {
+ AwsServiceException ex409 = S3Exception.builder().statusCode(409).build();
+ when(mockS3Client.getObject(any(GetObjectRequest.class))).thenThrow(ex409);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+ assertEquals(LockGetResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).info(contains("Conflicting operation has occurred"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testGetCurrentLockFile_objectFound() {
+ StorageLockData lockData = new StorageLockData(false, 9999L, "myTxOwner");
+ mockS3ObjectWithLockData(lockData, "abc-etag");
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+ assertEquals(LockGetResult.SUCCESS, result.getLeft());
+ assertTrue(result.getRight().isPresent());
+ assertEquals("abc-etag", result.getRight().get().getVersionId());
+ assertEquals("myTxOwner", result.getRight().get().getOwner());
+ }
+
+ @Test
+ void testGetCurrentLockFile_rateLimit() {
+ AwsServiceException ex429 = S3Exception.builder().statusCode(429).build();
+ when(mockS3Client.getObject(any(GetObjectRequest.class))).thenThrow(ex429);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+ assertEquals(LockGetResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).warn(contains("Rate limit exceeded"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testGetCurrentLockFile_serverError() {
+ AwsServiceException ex500 = S3Exception.builder().statusCode(500).build();
+ when(mockS3Client.getObject(any(GetObjectRequest.class))).thenThrow(ex500);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+ assertEquals(LockGetResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty());
+ verify(mockLogger).warn(contains("S3 internal server error"), eq(OWNER_ID), eq(LOCK_FILE_PATH), eq(ex500));
+ }
+
+ @Test
+ void testGetCurrentLockFile_unexpectedError() {
+ AwsServiceException ex400 = S3Exception.builder().statusCode(400).build();
+ when(mockS3Client.getObject(any(GetObjectRequest.class))).thenThrow(ex400);
+
+ assertThrows(S3Exception.class, () -> lockService.readCurrentLockFile());
+ }
+
+ @Test
+ void testGetCurrentLockFile_readFailure() {
+ GetObjectResponse getObjectResponse = GetObjectResponse.builder().eTag("some-etag").build();
+ @SuppressWarnings("unchecked")
+ InputStream mockInputStream = new InputStream() {
+ @Override
+ public int read() throws IOException {
+ throw new IOException("Read error");
+ }
+ };
+ ResponseInputStream mockStream = new ResponseInputStream(getObjectResponse, mockInputStream);
+ when(mockS3Client.getObject(any(GetObjectRequest.class))).thenReturn(mockStream);
+ assertThrows(HoodieIOException.class, () -> lockService.readCurrentLockFile());
+ }
+
+ @Test
+ void testClose() throws Exception {
+ lockService.close();
+ verify(mockS3Client).close();
+ }
+}
diff --git a/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageLockClientAuditOperations.java b/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageLockClientAuditOperations.java
new file mode 100644
index 000000000000..b92433806644
--- /dev/null
+++ b/hudi-aws/src/test/java/org/apache/hudi/aws/transaction/lock/TestS3StorageLockClientAuditOperations.java
@@ -0,0 +1,272 @@
+/*
+ * 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.hudi.aws.transaction.lock;
+
+import org.apache.hudi.common.util.Option;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.core.ResponseBytes;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.GetObjectRequest;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
+import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.services.s3.model.S3Exception;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Properties;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.contains;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for S3StorageLockClient audit operations (readObject and writeObject methods)
+ */
+public class TestS3StorageLockClientAuditOperations {
+
+ private S3Client mockS3Client;
+ private Logger mockLogger;
+ private S3StorageLockClient lockClient;
+
+ @BeforeEach
+ void setUp() {
+ mockS3Client = mock(S3Client.class);
+ mockLogger = mock(Logger.class);
+ String ownerId = "test-owner";
+ String lockFileUri = "s3://test-bucket/table/.hoodie/.locks/table_lock.json";
+ lockClient = new S3StorageLockClient(
+ ownerId,
+ lockFileUri,
+ new Properties(),
+ (bucket, props) -> mockS3Client,
+ mockLogger);
+ }
+
+ @Test
+ void testReadConfigWithCheckExistsFirstFileNotFound() {
+ String configPath = "s3://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // HEAD request returns 404
+ S3Exception notFoundException = (S3Exception) S3Exception.builder()
+ .statusCode(404)
+ .message("Not Found")
+ .build();
+ when(mockS3Client.headObject(any(HeadObjectRequest.class)))
+ .thenThrow(notFoundException);
+
+ Option<String> result = lockClient.readObject(configPath, true);
+
+ assertTrue(result.isEmpty());
+ // Should only call HEAD, not GET
+ verify(mockS3Client, times(1)).headObject(any(HeadObjectRequest.class));
+ verify(mockS3Client, never()).getObjectAsBytes(any(GetObjectRequest.class));
+ }
+
+ @Test
+ void testReadConfigWithCheckExistsFirstFileExists() {
+ String configPath = "s3://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+ String expectedContent = "{\"STORAGE_LP_AUDIT_SERVICE_ENABLED\": true}";
+
+ // HEAD request succeeds
+ HeadObjectResponse headResponse = HeadObjectResponse.builder().build();
+ when(mockS3Client.headObject(any(HeadObjectRequest.class)))
+ .thenReturn(headResponse);
+
+ // GET request returns content
+ ResponseBytes<GetObjectResponse> responseBytes = ResponseBytes.fromByteArray(
+ GetObjectResponse.builder().build(),
+ expectedContent.getBytes(StandardCharsets.UTF_8));
+ when(mockS3Client.getObjectAsBytes(any(GetObjectRequest.class)))
+ .thenReturn(responseBytes);
+
+ Option<String> result = lockClient.readObject(configPath, true);
+
+ assertTrue(result.isPresent());
+ assertEquals(expectedContent, result.get());
+ // Should call both HEAD and GET
+ verify(mockS3Client, times(1)).headObject(any(HeadObjectRequest.class));
+ verify(mockS3Client, times(1)).getObjectAsBytes(any(GetObjectRequest.class));
+ }
+
+ @Test
+ void testReadConfigWithoutCheckExistsFirstFileNotFound() {
+ String configPath = "s3://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // Direct GET request returns 404
+ S3Exception notFoundException = (S3Exception) S3Exception.builder()
+ .statusCode(404)
+ .message("Not Found")
+ .build();
+ when(mockS3Client.getObjectAsBytes(any(GetObjectRequest.class)))
+ .thenThrow(notFoundException);
+
+ Option<String> result = lockClient.readObject(configPath, false);
+
+ assertTrue(result.isEmpty());
+ // Should not call HEAD, only GET
+ verify(mockS3Client, never()).headObject(any(HeadObjectRequest.class));
+ verify(mockS3Client, times(1)).getObjectAsBytes(any(GetObjectRequest.class));
+ }
+
+ @Test
+ void testReadConfigWithoutCheckExistsFirstFileExists() {
+ String configPath = "s3://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+ String expectedContent = "{\"STORAGE_LP_AUDIT_SERVICE_ENABLED\": false}";
+
+ // Direct GET request returns content
+ ResponseBytes<GetObjectResponse> responseBytes = ResponseBytes.fromByteArray(
+ GetObjectResponse.builder().build(),
+ expectedContent.getBytes(StandardCharsets.UTF_8));
+ when(mockS3Client.getObjectAsBytes(any(GetObjectRequest.class)))
+ .thenReturn(responseBytes);
+
+ Option<String> result = lockClient.readObject(configPath, false);
+
+ assertTrue(result.isPresent());
+ assertEquals(expectedContent, result.get());
+ // Should not call HEAD, only GET
+ verify(mockS3Client, never()).headObject(any(HeadObjectRequest.class));
+ verify(mockS3Client, times(1)).getObjectAsBytes(any(GetObjectRequest.class));
+ }
+
+ @Test
+ void testReadConfigWithCheckExistsFirstOtherS3Error() {
+ String configPath = "s3://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // HEAD request returns non-404 error
+ S3Exception serverError = (S3Exception) S3Exception.builder()
+ .statusCode(500)
+ .message("Internal Server Error")
+ .build();
+ when(mockS3Client.headObject(any(HeadObjectRequest.class)))
+ .thenThrow(serverError);
+
+ Option<String> result = lockClient.readObject(configPath, true);
+
+ assertTrue(result.isEmpty());
+ verify(mockS3Client, times(1)).headObject(any(HeadObjectRequest.class));
+ verify(mockS3Client, never()).getObjectAsBytes(any(GetObjectRequest.class));
+ }
+
+ @Test
+ void testReadConfigWithInvalidUri() {
+ String invalidPath = "not-a-valid-uri";
+
+ Option<String> result = lockClient.readObject(invalidPath, false);
+
+ assertTrue(result.isEmpty());
+ // Should not make any S3 calls due to URI parsing error
+ verify(mockS3Client, never()).headObject(any(HeadObjectRequest.class));
+ }
+
+ @Test
+ void testReadConfigWithRateLimitError() {
+ String configPath = "s3://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // GET request returns rate limit error
+ S3Exception rateLimitException = (S3Exception) S3Exception.builder()
+ .statusCode(429)
+ .message("Too Many Requests")
+ .build();
+ when(mockS3Client.getObjectAsBytes(any(GetObjectRequest.class)))
+ .thenThrow(rateLimitException);
+
+ Option<String> result = lockClient.readObject(configPath, false);
+
+ assertTrue(result.isEmpty());
+ verify(mockS3Client, times(1)).getObjectAsBytes(any(GetObjectRequest.class));
+ }
+
+ // ================================
+ // writeObject() tests
+ // ================================
+
+ @Test
+ void testWriteObject_success() {
+ String filePath = "s3://test-bucket/audit/test-audit.jsonl";
+ String content = "{\"test\": \"data\"}\n";
+ PutObjectResponse putResp = PutObjectResponse.builder().eTag("write-etag-123").build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenReturn(putResp);
+
+ boolean result = lockClient.writeObject(filePath, content);
+
+ assertTrue(result);
+ verify(mockS3Client, times(1)).putObject(
+ eq(PutObjectRequest.builder().bucket("test-bucket").key("audit/test-audit.jsonl").build()),
+ any(RequestBody.class)
+ );
+ verify(mockLogger).debug("Successfully wrote object to: {}", filePath);
+ }
+
+ @Test
+ void testWriteObject_s3Exception() {
+ String filePath = "s3://test-bucket/audit/test-audit.jsonl";
+ String content = "{\"test\": \"data\"}\n";
+ AwsServiceException s3Exception = S3Exception.builder().statusCode(500).build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenThrow(s3Exception);
+
+ boolean result = lockClient.writeObject(filePath, content);
+
+ assertFalse(result);
+ verify(mockLogger).warn(contains("Error writing object to"), eq(filePath), eq(s3Exception));
+ }
+
+ @Test
+ void testWriteObject_invalidPath() {
+ String invalidPath = "invalid-path";
+ String content = "{\"test\": \"data\"}\n";
+
+ boolean result = lockClient.writeObject(invalidPath, content);
+
+ assertFalse(result);
+ verify(mockLogger).warn(contains("Error writing object to"), eq(invalidPath), any(Exception.class));
+ }
+
+ @Test
+ void testWriteObject_emptyContent() {
+ String filePath = "s3://test-bucket/audit/empty-content.jsonl";
+ String content = "";
+ PutObjectResponse putResp = PutObjectResponse.builder().eTag("empty-etag-456").build();
+ when(mockS3Client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenReturn(putResp);
+
+ boolean result = lockClient.writeObject(filePath, content);
+
+ assertTrue(result);
+ verify(mockS3Client, times(1)).putObject(
+ eq(PutObjectRequest.builder().bucket("test-bucket").key("audit/empty-content.jsonl").build()),
+ any(RequestBody.class)
+ );
+ verify(mockLogger).debug("Successfully wrote object to: {}", filePath);
+ }
+}
diff --git a/hudi-cli/src/main/java/org/apache/hudi/cli/commands/LockAuditingCommand.java b/hudi-cli/src/main/java/org/apache/hudi/cli/commands/LockAuditingCommand.java
new file mode 100644
index 000000000000..8545ad942b5a
--- /dev/null
+++ b/hudi-cli/src/main/java/org/apache/hudi/cli/commands/LockAuditingCommand.java
@@ -0,0 +1,661 @@
+/*
+ * 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.hudi.cli.commands;
+
+import org.apache.hudi.cli.HoodieCLI;
+import org.apache.hudi.client.transaction.lock.audit.StorageLockProviderAuditService;
+import org.apache.hudi.common.util.FileIOUtils;
+import org.apache.hudi.common.util.Option;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.Path;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.shell.standard.ShellComponent;
+import org.springframework.shell.standard.ShellMethod;
+import org.springframework.shell.standard.ShellOption;
+
+import java.io.BufferedReader;
+import java.io.FileNotFoundException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+
+/**
+ * CLI commands for managing Hudi table lock auditing functionality.
+ */
+@ShellComponent
+public class LockAuditingCommand {
+
+ private static final Logger LOG = LoggerFactory.getLogger(LockAuditingCommand.class);
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ /**
+ * Represents a single audit log entry in the JSONL file.
+ * Maps to the structure written by StorageLockProviderAuditService.
+ */
+ public static class AuditRecord {
+ private final String ownerId;
+ private final long transactionStartTime;
+ private final long timestamp;
+ private final String state;
+ private final long lockExpiration;
+ private final boolean lockHeld;
+
+ @JsonCreator
+ public AuditRecord(
+ @JsonProperty("ownerId") String ownerId,
+ @JsonProperty("transactionStartTime") long transactionStartTime,
+ @JsonProperty("timestamp") long timestamp,
+ @JsonProperty("state") String state,
+ @JsonProperty("lockExpiration") long lockExpiration,
+ @JsonProperty("lockHeld") boolean lockHeld) {
+ this.ownerId = ownerId;
+ this.transactionStartTime = transactionStartTime;
+ this.timestamp = timestamp;
+ this.state = state;
+ this.lockExpiration = lockExpiration;
+ this.lockHeld = lockHeld;
+ }
+
+ public String getOwnerId() {
+ return ownerId;
+ }
+
+ public long getTransactionStartTime() {
+ return transactionStartTime;
+ }
+
+ public long getTimestamp() {
+ return timestamp;
+ }
+
+ public String getState() {
+ return state;
+ }
+
+ public long getLockExpiration() {
+ return lockExpiration;
+ }
+
+ public boolean isLockHeld() {
+ return lockHeld;
+ }
+ }
+
+ /**
+ * Represents a transaction window with start time, end time, and metadata.
+ */
+ static class TransactionWindow {
+ final String ownerId;
+ final long transactionCreationTime;
+ final long lockAcquisitionTime;
+ final Option<Long> lockReleaseTime;
+ final Option<Long> lockExpirationTime;
+ final String filename;
+
+ TransactionWindow(String ownerId, long transactionCreationTime, long lockAcquisitionTime,
+ Option<Long> lockReleaseTime, Option<Long> lockExpirationTime, String filename) {
+ this.ownerId = ownerId;
+ this.transactionCreationTime = transactionCreationTime;
+ this.lockAcquisitionTime = lockAcquisitionTime;
+ this.lockReleaseTime = lockReleaseTime;
+ this.lockExpirationTime = lockExpirationTime;
+ this.filename = filename;
+ }
+
+ long getEffectiveEndTime() {
+ return lockReleaseTime.orElse(lockExpirationTime.orElse(lockAcquisitionTime));
+ }
+ }
+
+ /**
+ * Holds validation results.
+ */
+ static class ValidationResults {
+ final List<String> errors;
+ final List<String> warnings;
+
+ ValidationResults(List<String> errors, List<String> warnings) {
+ this.errors = errors;
+ this.warnings = warnings;
+ }
+ }
+
+ /**
+ * Holds cleanup operation results.
+ */
+ static class CleanupResult {
+ final boolean success;
+ final int deletedCount;
+ final int failedCount;
+ final int totalFiles;
+ final String message;
+ final boolean isDryRun;
+
+ CleanupResult(boolean success, int deletedCount, int failedCount, int totalFiles, String message, boolean isDryRun) {
+ this.success = success;
+ this.deletedCount = deletedCount;
+ this.failedCount = failedCount;
+ this.totalFiles = totalFiles;
+ this.message = message;
+ this.isDryRun = isDryRun;
+ }
+
+ boolean hasErrors() {
+ return failedCount > 0;
+ }
+
+ boolean hasFiles() {
+ return totalFiles > 0;
+ }
+
+ @Override
+ public String toString() {
+ return message;
+ }
+ }
+
+ /**
+ * Enables lock audit logging for the currently connected Hudi table.
+ * This command creates or updates the audit configuration file to enable
+ * audit logging for storage lock operations.
+ *
+ * @return Status message indicating success or failure
+ */
+ @ShellMethod(key = "locks audit enable", value = "Enable storage lock audit service for the current table")
+ public String enableLockAudit() {
+
+ if (HoodieCLI.basePath == null) {
+ return "No Hudi table loaded. Please connect to a table first.";
+ }
+
+ try {
+ // Create the audit config file path using utility method
+ String auditConfigPath = StorageLockProviderAuditService.getAuditConfigPath(HoodieCLI.basePath);
+
+ // Create the JSON content
+ ObjectNode configJson = OBJECT_MAPPER.createObjectNode();
+ configJson.put(StorageLockProviderAuditService.STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD, true);
+ String jsonContent = OBJECT_MAPPER.writeValueAsString(configJson);
+
+ // Write the config file using HoodieStorage
+ Path configPath = new Path(auditConfigPath);
+ try (OutputStream outputStream = HoodieCLI.fs.create(configPath, true)) {
+ outputStream.write(jsonContent.getBytes());
+ }
+
+ return String.format("Lock audit enabled successfully.\nAudit config written to: %s\n"
+ + "Audit files will be stored at: %s", auditConfigPath, StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath));
+
+ } catch (Exception e) {
+ LOG.error("Error enabling lock audit", e);
+ return String.format("Failed to enable lock audit: %s", e.getMessage());
+ }
+ }
+
+ /**
+ * Disables lock audit logging for the currently connected Hudi table.
+ * This command updates the audit configuration file to disable audit logging.
+ *
+ * @param keepAuditFiles Whether to preserve existing audit files when disabling
+ * @return Status message indicating success or failure
+ */
+ @ShellMethod(key = "locks audit disable", value = "Disable storage lock audit service for the current table")
+ public String disableLockAudit(
+ @ShellOption(value = {"--keepAuditFiles"}, defaultValue = "true",
+ help = "Keep existing audit files when disabling") final boolean keepAuditFiles) {
+
+ if (HoodieCLI.basePath == null) {
+ return "No Hudi table loaded. Please connect to a table first.";
+ }
+
+ try {
+ // Create the audit config file path
+ String auditConfigPath = StorageLockProviderAuditService.getAuditConfigPath(HoodieCLI.basePath);
+ Path configPath = new Path(auditConfigPath);
+
+ // Check if config file exists by attempting to get its info
+ if (!HoodieCLI.fs.exists(configPath)) {
+ return "Lock audit is already disabled (no configuration file found).";
+ }
+
+ // Create the JSON content with audit disabled
+ ObjectNode configJson = OBJECT_MAPPER.createObjectNode();
+ configJson.put(StorageLockProviderAuditService.STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD, false);
+ String jsonContent = OBJECT_MAPPER.writeValueAsString(configJson);
+
+ // Write the config file
+ try (OutputStream outputStream = HoodieCLI.fs.create(configPath, true)) {
+ outputStream.write(jsonContent.getBytes());
+ }
+
+ String message = String.format("Lock audit disabled successfully.\nAudit config updated at: %s", auditConfigPath);
+
+ if (keepAuditFiles) {
+ message += String.format("\nExisting audit files preserved at: %s", StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath));
+ } else {
+ // Call the cleanup method to prune all old files
+ CleanupResult cleanupResult = performAuditCleanup(false, 0);
+ if (cleanupResult.success && cleanupResult.deletedCount > 0) {
+ message += String.format("\nAudit files cleaned up at: %s", StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath));
+ } else if (cleanupResult.success && cleanupResult.deletedCount == 0) {
+ message += String.format("\nNo audit files to clean up at: %s", StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath));
+ } else {
+ message += String.format("\nWarning: Some audit files may not have been cleaned up: %s", cleanupResult.message);
+ }
+ }
+
+ return message;
+
+ } catch (Exception e) {
+ LOG.error("Error disabling lock audit", e);
+ return String.format("Failed to disable lock audit: %s", e.getMessage());
+ }
+ }
+
+ /**
+ * Shows the current status of lock audit logging for the connected table.
+ * This command checks the audit configuration file and reports whether
+ * auditing is currently enabled or disabled.
+ *
+ * @return Status information about the current audit configuration
+ */
+ @ShellMethod(key = "locks audit status", value = "Show the current status of lock audit service")
+ public String showLockAuditStatus() {
+
+ if (HoodieCLI.basePath == null) {
+ return "No Hudi table loaded. Please connect to a table first.";
+ }
+
+ try {
+ // Create the audit config file path
+ String auditConfigPath = StorageLockProviderAuditService.getAuditConfigPath(HoodieCLI.basePath);
+ Path configPath = new Path(auditConfigPath);
+
+ // Read and parse the configuration
+ String configContent;
+ try (InputStream inputStream = HoodieCLI.fs.open(configPath)) {
+ configContent = new String(FileIOUtils.readAsByteArray(inputStream));
+ } catch (FileNotFoundException e) {
+ return String.format("Lock Audit Status: DISABLED\n"
+ + "Table: %s\n"
+ + "Config file: %s (not found)\n"
+ + "Use 'locks audit enable' to enable audit logging.",
+ HoodieCLI.basePath, auditConfigPath);
+ }
+ JsonNode rootNode = OBJECT_MAPPER.readTree(configContent);
+ JsonNode enabledNode = rootNode.get(StorageLockProviderAuditService.STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD);
+ boolean isEnabled = enabledNode != null && enabledNode.asBoolean(false);
+
+ String status = isEnabled ? "ENABLED" : "DISABLED";
+
+ return String.format("Lock Audit Status: %s\n"
+ + "Table: %s\n"
+ + "Config file: %s\n"
+ + "Audit files location: %s",
+ status, HoodieCLI.basePath, auditConfigPath, StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath));
+
+ } catch (Exception e) {
+ LOG.error("Error checking lock audit status", e);
+ return String.format("Failed to check lock audit status: %s", e.getMessage());
+ }
+ }
+
+ /**
+ * Validates the audit lock files for consistency and integrity.
+ * This command checks for issues such as corrupted files, invalid format,
+ * incorrect state transitions, and orphaned locks.
+ *
+ * @return Validation results including any issues found
+ */
+ @ShellMethod(key = "locks audit validate", value = "Validate audit lock files for consistency and integrity")
+ public String validateAuditLocks() {
+
+ if (HoodieCLI.basePath == null) {
+ return "No Hudi table loaded. Please connect to a table first.";
+ }
+
+ try {
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditFolder = new Path(auditFolderPath);
+
+ // Get all audit files
+ FileStatus[] allFiles;
+ try {
+ allFiles = HoodieCLI.fs.listStatus(auditFolder);
+ } catch (FileNotFoundException e) {
+ return "Validation Result: PASSED\n"
+ + "Transactions Validated: 0\n"
+ + "Issues Found: 0\n"
+ + "Details: No audit folder found - nothing to validate";
+ }
+ List<FileStatus> auditFiles = new ArrayList<>();
+ for (FileStatus pathInfo : allFiles) {
+ if (pathInfo.isFile() && pathInfo.getPath().getName().endsWith(".jsonl")) {
+ auditFiles.add(pathInfo);
+ }
+ }
+
+ if (auditFiles.isEmpty()) {
+ return "Validation Result: PASSED\n"
+ + "Transactions Validated: 0\n"
+ + "Issues Found: 0\n"
+ + "Details: No audit files found - nothing to validate";
+ }
+
+ // Parse all audit files into transaction windows
+ List<TransactionWindow> windows = new ArrayList<>();
+ List<String> parseErrors = new ArrayList<>();
+ for (FileStatus pathInfo : auditFiles) {
+ Option<TransactionWindow> window = parseAuditFile(pathInfo);
+ if (window.isPresent()) {
+ windows.add(window.get());
+ } else {
+ parseErrors.add(String.format("[ERROR] Failed to parse audit file: %s", pathInfo.getPath().getName()));
+ }
+ }
+
+ if (windows.isEmpty()) {
+ return String.format("Validation Result: FAILED\n"
+ + "Transactions Validated: 0\n"
+ + "Issues Found: %d\n"
+ + "Details: Failed to parse any audit files", auditFiles.size());
+ }
+
+ // Validate transactions
+ ValidationResults validationResults = validateTransactionWindows(windows);
+
+ // Add parse errors to the validation results
+ validationResults.errors.addAll(parseErrors);
+
+ // Generate result
+ int totalIssues = validationResults.errors.size() + validationResults.warnings.size();
+ String result;
+ String details;
+
+ if (totalIssues == 0) {
+ result = "PASSED";
+ details = "All audit lock transactions validated successfully";
+ } else {
+ result = validationResults.errors.isEmpty() ? "WARNING" : "FAILED";
+ List<String> allIssues = new ArrayList<>();
+ allIssues.addAll(validationResults.errors);
+ allIssues.addAll(validationResults.warnings);
+ details = String.join(", ", allIssues);
+ }
+
+ return String.format("Validation Result: %s\n"
+ + "Audit Files: %d total, %d parsed successfully, %d failed to parse\n"
+ + "Transactions Validated: %d\n"
+ + "Issues Found: %d\n"
+ + "Details: %s", result, auditFiles.size(), windows.size(), parseErrors.size(), windows.size(), totalIssues, details);
+
+ } catch (Exception e) {
+ LOG.error("Error validating audit locks", e);
+ return String.format("Validation Result: ERROR\n"
+ + "Transactions Validated: 0\n"
+ + "Issues Found: -1\n"
+ + "Details: Validation failed: %s", e.getMessage());
+ }
+ }
+
+ /**
+ * Cleans up old audit lock files based on age threshold.
+ * This command removes audit files that are older than the specified number of days.
+ *
+ * @param dryRun Whether to perform a dry run (preview changes without deletion)
+ * @param ageDays Number of days to keep audit files (default 7)
+ * @return Status message indicating files cleaned or to be cleaned
+ */
+ @ShellMethod(key = "locks audit cleanup", value = "Clean up old audit lock files")
+ public String cleanupAuditLocks(
+ @ShellOption(value = {"--dryRun"}, defaultValue = "false",
+ help = "Preview changes without actually deleting files") final boolean dryRun,
+ @ShellOption(value = {"--ageDays"}, defaultValue = "7",
+ help = "Delete audit files older than this many days") final String ageDaysStr) {
+
+ try {
+ if (HoodieCLI.basePath == null) {
+ return "No Hudi table loaded. Please connect to a table first.";
+ }
+
+ // Parse ageDays manually to handle validation properly
+ int ageDays;
+ try {
+ ageDays = Integer.parseInt(ageDaysStr);
+ } catch (NumberFormatException e) {
+ return "Error: ageDays must be a value greater than 0.";
+ }
+
+ if (ageDays < 0) {
+ return "Error: ageDays must be non-negative (>= 0).";
+ }
+
+ return performAuditCleanup(dryRun, ageDays).toString();
+ } catch (Exception e) {
+ LOG.error("Error during audit cleanup", e);
+ return String.format("Error during cleanup: %s", e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName());
+ }
+ }
+
+ /**
+ * Internal method to perform audit cleanup. Used by both the CLI command and disable method.
+ *
+ * @param dryRun Whether to perform a dry run (preview changes without deletion)
+ * @param ageDays Number of days to keep audit files (0 means delete all, must be non-negative)
+ * @return CleanupResult containing cleanup operation details
+ */
+ private CleanupResult performAuditCleanup(boolean dryRun, int ageDays) {
+ try {
+ if (HoodieCLI.fs == null) {
+ String message = "Storage not initialized.";
+ return new CleanupResult(false, 0, 0, 0, message, dryRun);
+ }
+
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditFolder = new Path(auditFolderPath);
+
+ // Calculate cutoff timestamp (ageDays ago)
+ long cutoffTime = System.currentTimeMillis() - (ageDays * 24L * 60L * 60L * 1000L);
+
+ // List all files in audit folder and filter by modification time
+ FileStatus[] allFiles;
+ try {
+ allFiles = HoodieCLI.fs.listStatus(auditFolder);
+ } catch (FileNotFoundException e) {
+ String message = "No audit folder found - nothing to cleanup.";
+ return new CleanupResult(true, 0, 0, 0, message, dryRun);
+ }
+ List<FileStatus> auditFiles = new ArrayList<>();
+ List<FileStatus> oldFiles = new ArrayList<>();
+
+ // Filter to get only .jsonl files
+ for (FileStatus pathInfo : allFiles) {
+ if (pathInfo.isFile() && pathInfo.getPath().getName().endsWith(".jsonl")) {
+ auditFiles.add(pathInfo);
+ if (pathInfo.getModificationTime() < cutoffTime) {
+ oldFiles.add(pathInfo);
+ }
+ }
+ }
+
+ if (oldFiles.isEmpty()) {
+ String message = String.format("No audit files older than %d days found.", ageDays);
+ return new CleanupResult(true, 0, 0, auditFiles.size(), message, dryRun);
+ }
+
+ int fileCount = oldFiles.size();
+
+ if (dryRun) {
+ String message = String.format("Dry run: Would delete %d audit files older than %d days.", fileCount, ageDays);
+ return new CleanupResult(true, 0, 0, fileCount, message, dryRun);
+ } else {
+ // Actually delete the files
+ int deletedCount = 0;
+ int failedCount = 0;
+
+ for (FileStatus pathInfo : oldFiles) {
+ try {
+ HoodieCLI.fs.delete(pathInfo.getPath(), false);
+ deletedCount++;
+ } catch (Exception e) {
+ failedCount++;
+ LOG.warn("Failed to delete audit file: " + pathInfo.getPath(), e);
+ }
+ }
+
+ if (failedCount == 0) {
+ String message = String.format("Successfully deleted %d audit files older than %d days.", deletedCount, ageDays);
+ return new CleanupResult(true, deletedCount, failedCount, fileCount, message, dryRun);
+ } else {
+ String message = String.format("Deleted %d audit files, failed to delete %d files.", deletedCount, failedCount);
+ return new CleanupResult(false, deletedCount, failedCount, fileCount, message, dryRun);
+ }
+ }
+
+ } catch (Exception e) {
+ LOG.error("Error cleaning up audit locks", e);
+ String message = String.format("Failed to cleanup audit locks: %s", e.getMessage());
+ return new CleanupResult(false, 0, 0, 0, message, dryRun);
+ }
+ }
+
+ /**
+ * Parses an audit file and extracts transaction window information.
+ */
+ private Option<TransactionWindow> parseAuditFile(FileStatus pathInfo) {
+ String filename = pathInfo.getPath().getName();
+
+ try {
+ // Read and parse JSONL content
+ List<AuditRecord> entries = new ArrayList<>();
+ try (InputStream inputStream = HoodieCLI.fs.open(pathInfo.getPath());
+ BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
+ String line;
+ while ((line = reader.readLine()) != null) {
+ if (line.trim().isEmpty()) {
+ continue;
+ }
+ try {
+ AuditRecord entry = OBJECT_MAPPER.readValue(line, AuditRecord.class);
+ entries.add(entry);
+ } catch (Exception e) {
+ LOG.warn("Failed to parse JSON line in file " + filename + ": " + line, e);
+ }
+ }
+ }
+
+ if (entries.isEmpty()) {
+ return Option.empty();
+ }
+
+ // Extract transaction metadata from first entry
+ AuditRecord firstEntry = entries.get(0);
+ String ownerId = firstEntry.getOwnerId();
+ long transactionCreationTime = firstEntry.getTransactionStartTime();
+
+ // Find first START timestamp
+ long lockAcquisitionTime = transactionCreationTime; // default to transaction creation time
+ for (AuditRecord entry : entries) {
+ if ("START".equals(entry.getState())) {
+ lockAcquisitionTime = entry.getTimestamp();
+ break;
+ }
+ }
+
+ // Find last END timestamp
+ Option<Long> lockReleaseTime = Option.empty();
+ for (int i = entries.size() - 1; i >= 0; i--) {
+ AuditRecord entry = entries.get(i);
+ if ("END".equals(entry.getState())) {
+ lockReleaseTime = Option.of(entry.getTimestamp());
+ break;
+ }
+ }
+
+ // Find last expiration time as fallback
+ Option<Long> lockExpirationTime = Option.empty();
+ if (!entries.isEmpty()) {
+ AuditRecord lastEntry = entries.get(entries.size() - 1);
+ lockExpirationTime = Option.of(lastEntry.getLockExpiration());
+ }
+
+ return Option.of(new TransactionWindow(
+ ownerId,
+ transactionCreationTime,
+ lockAcquisitionTime,
+ lockReleaseTime,
+ lockExpirationTime,
+ filename
+ ));
+ } catch (Exception e) {
+ LOG.warn("Failed to parse audit file: " + filename, e);
+ return Option.empty();
+ }
+ }
+
+ /**
+ * Validates transaction windows for overlaps and proper closure.
+ */
+ private ValidationResults validateTransactionWindows(List<TransactionWindow> windows) {
+ List<String> errors = new ArrayList<>();
+ List<String> warnings = new ArrayList<>();
+
+ // Check for transactions without proper END
+ for (TransactionWindow window : windows) {
+ if (!window.lockReleaseTime.isPresent()) {
+ warnings.add(String.format("[WARNING] %s => transaction did not end gracefully. This could be due to driver OOM or non-graceful shutdown.",
+ window.filename));
+ }
+ }
+
+ // Sort windows by start time for overlap detection
+ List<TransactionWindow> sortedWindows = new ArrayList<>(windows);
+ sortedWindows.sort(Comparator.comparingLong(w -> w.lockAcquisitionTime));
+
+ // Check for overlaps
+ for (int i = 0; i < sortedWindows.size(); i++) {
+ TransactionWindow currentWindow = sortedWindows.get(i);
+ long currentEnd = currentWindow.getEffectiveEndTime();
+
+ // Check all subsequent windows for overlaps
+ for (int j = i + 1; j < sortedWindows.size(); j++) {
+ TransactionWindow otherWindow = sortedWindows.get(j);
+ long otherStart = otherWindow.lockAcquisitionTime;
+
+ // Check if windows overlap
+ if (otherStart < currentEnd) {
+ errors.add(String.format("[ERROR] %s => overlaps with %s",
+ currentWindow.filename, otherWindow.filename));
+ }
+ }
+ }
+
+ return new ValidationResults(errors, warnings);
+ }
+}
\ No newline at end of file
diff --git a/hudi-cli/src/test/java/org/apache/hudi/cli/commands/TestLockAuditingCommand.java b/hudi-cli/src/test/java/org/apache/hudi/cli/commands/TestLockAuditingCommand.java
new file mode 100644
index 000000000000..c2cc6e0d63b5
--- /dev/null
+++ b/hudi-cli/src/test/java/org/apache/hudi/cli/commands/TestLockAuditingCommand.java
@@ -0,0 +1,945 @@
+/*
+ * 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.hudi.cli.commands;
+
+import org.apache.hudi.cli.HoodieCLI;
+import org.apache.hudi.cli.functional.CLIFunctionalTestHarness;
+import org.apache.hudi.cli.testutils.ShellEvaluationResultUtil;
+import org.apache.hudi.client.transaction.lock.audit.StorageLockProviderAuditService;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.common.model.HoodieTableType;
+import org.apache.hudi.common.util.FileIOUtils;
+
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.Path;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.BeforeEach;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.shell.Shell;
+
+import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Test class for {@link org.apache.hudi.cli.commands.LockAuditingCommand}.
+ */
+@Tag("functional")
+@SpringBootTest(properties = {"spring.shell.interactive.enabled=false", "spring.shell.command.script.enabled=false"})
+public class TestLockAuditingCommand extends CLIFunctionalTestHarness {
+
+ @Autowired
+ private Shell shell;
+
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ /**
+ * Helper method to create AuditRecord for tests with lockHeld defaulting to true.
+ */
+ private static LockAuditingCommand.AuditRecord createAuditRecord(
+ String ownerId, long transactionStartTime, long timestamp, String state, long lockExpiration) {
+ return new LockAuditingCommand.AuditRecord(ownerId, transactionStartTime, timestamp, state, lockExpiration, true);
+ }
+
+ /**
+ * Represents a transaction scenario with its audit records
+ */
+ static class TransactionScenario {
+ final String filename; // e.g., "1234567890_owner1.jsonl"
+ final List<LockAuditingCommand.AuditRecord> records;
+
+ TransactionScenario(String filename, List<LockAuditingCommand.AuditRecord> records) {
+ this.filename = filename;
+ this.records = records;
+ }
+ }
+
+ /**
+ * Helper method to create audit files from scenarios
+ */
+ private void createAuditFiles(List<TransactionScenario> scenarios) throws IOException {
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditDir = new Path(auditFolderPath);
+
+ // Create audit directory if it doesn't exist
+ if (!HoodieCLI.fs.exists(auditDir)) {
+ HoodieCLI.fs.mkdirs(auditDir);
+ }
+
+ for (TransactionScenario scenario : scenarios) {
+ Path filePath = new Path(auditDir, scenario.filename);
+ StringBuilder jsonLines = new StringBuilder();
+ for (LockAuditingCommand.AuditRecord record : scenario.records) {
+ if (jsonLines.length() > 0) {
+ jsonLines.append("\n");
+ }
+ jsonLines.append(OBJECT_MAPPER.writeValueAsString(record));
+ }
+
+ try (OutputStream outputStream = HoodieCLI.fs.create(filePath, true)) {
+ outputStream.write(jsonLines.toString().getBytes());
+ }
+ }
+ }
+
+ @BeforeEach
+ public void setUp() throws Exception {
+ HoodieCLI.conf = jsc().hadoopConfiguration();
+ String tableName = tableName();
+ String tablePath = tablePath(tableName);
+ HoodieCLI.basePath = tablePath;
+
+ // Initialize table
+ HoodieTableMetaClient.withPropertyBuilder()
+ .setTableType(HoodieTableType.COPY_ON_WRITE.name())
+ .setTableName(tableName)
+ .setRecordKeyFields("key")
+ .initTable(HoodieCLI.conf, tablePath);
+
+ // Initialize storage
+ HoodieCLI.initFS(true);
+ }
+
+ /**
+ * Test enabling lock audit when no table is loaded.
+ */
+ @Test
+ public void testEnableLockAuditNoTable() {
+ // Clear the base path to simulate no table loaded
+ String originalBasePath = HoodieCLI.basePath;
+ HoodieCLI.basePath = null;
+
+ try {
+ Object result = shell.evaluate(() -> "locks audit enable");
+ assertAll("Command runs with no table loaded",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertEquals("No Hudi table loaded. Please connect to a table first.", result.toString()));
+ } finally {
+ HoodieCLI.basePath = originalBasePath;
+ }
+ }
+
+ /**
+ * Test disabling lock audit when no table is loaded.
+ */
+ @Test
+ public void testDisableLockAuditNoTable() {
+ // Clear the base path to simulate no table loaded
+ String originalBasePath = HoodieCLI.basePath;
+ HoodieCLI.basePath = null;
+
+ try {
+ Object result = shell.evaluate(() -> "locks audit disable");
+ assertAll("Command runs with no table loaded",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertEquals("No Hudi table loaded. Please connect to a table first.", result.toString()));
+ } finally {
+ HoodieCLI.basePath = originalBasePath;
+ }
+ }
+
+ /**
+ * Test showing lock audit status when no table is loaded.
+ */
+ @Test
+ public void testShowLockAuditStatusNoTable() {
+ // Clear the base path to simulate no table loaded
+ String originalBasePath = HoodieCLI.basePath;
+ HoodieCLI.basePath = null;
+
+ try {
+ Object result = shell.evaluate(() -> "locks audit status");
+ assertAll("Command runs with no table loaded",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertEquals("No Hudi table loaded. Please connect to a table first.", result.toString()));
+ } finally {
+ HoodieCLI.basePath = originalBasePath;
+ }
+ }
+
+ /**
+ * Test enabling lock audit successfully.
+ */
+ @Test
+ public void testEnableLockAuditSuccess() throws Exception {
+ Object result = shell.evaluate(() -> "locks audit enable");
+
+ assertAll("Enable command runs successfully",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Lock audit enabled successfully")));
+
+ // Verify the config file was created with correct content
+ String lockFolderPath = String.format("%s%s.hoodie%s.locks", HoodieCLI.basePath, Path.SEPARATOR, Path.SEPARATOR);
+ String auditConfigPath = String.format("%s%s%s", lockFolderPath, Path.SEPARATOR, StorageLockProviderAuditService.AUDIT_CONFIG_FILE_NAME);
+ Path configPath = new Path(auditConfigPath);
+
+ assertTrue(HoodieCLI.fs.exists(configPath), "Config file should exist");
+
+ String configContent;
+ try (InputStream inputStream = HoodieCLI.fs.open(configPath)) {
+ configContent = new String(FileIOUtils.readAsByteArray(inputStream));
+ }
+ JsonNode rootNode = OBJECT_MAPPER.readTree(configContent);
+ JsonNode enabledNode = rootNode.get(StorageLockProviderAuditService.STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD);
+
+ assertNotNull(enabledNode, "Config should contain enabled field");
+ assertTrue(enabledNode.asBoolean(), "Audit should be enabled");
+ }
+
+ /**
+ * Test showing lock audit status when disabled (no config file).
+ */
+ @Test
+ public void testShowLockAuditStatusDisabled() {
+ Object result = shell.evaluate(() -> "locks audit status");
+
+ assertAll("Status command shows disabled state",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Lock Audit Status: DISABLED")),
+ () -> assertTrue(result.toString().contains("(not found)")));
+ }
+
+ /**
+ * Test the complete workflow: enable, check status, disable, check status.
+ */
+ @Test
+ public void testCompleteAuditWorkflow() throws Exception {
+ // 1. Enable audit
+ Object enableResult = shell.evaluate(() -> "locks audit enable");
+ assertTrue(ShellEvaluationResultUtil.isSuccess(enableResult));
+ assertTrue(enableResult.toString().contains("Lock audit enabled successfully"));
+
+ // 2. Check status - should be enabled
+ Object statusEnabledResult = shell.evaluate(() -> "locks audit status");
+ assertTrue(ShellEvaluationResultUtil.isSuccess(statusEnabledResult));
+ assertTrue(statusEnabledResult.toString().contains("Lock Audit Status: ENABLED"));
+
+ // 3. Disable audit with default keepAuditFiles=true
+ Object disableResult = shell.evaluate(() -> "locks audit disable");
+ assertTrue(ShellEvaluationResultUtil.isSuccess(disableResult));
+ assertTrue(disableResult.toString().contains("Lock audit disabled successfully"));
+ assertTrue(disableResult.toString().contains("Existing audit files preserved"));
+
+ // 4. Check status - should be disabled
+ Object statusDisabledResult = shell.evaluate(() -> "locks audit status");
+ assertTrue(ShellEvaluationResultUtil.isSuccess(statusDisabledResult));
+ assertTrue(statusDisabledResult.toString().contains("Lock Audit Status: DISABLED"));
+
+ // Verify the config file still exists but with audit disabled
+ String lockFolderPath = String.format("%s%s.hoodie%s.locks", HoodieCLI.basePath, Path.SEPARATOR, Path.SEPARATOR);
+ String auditConfigPath = String.format("%s%s%s", lockFolderPath, Path.SEPARATOR, StorageLockProviderAuditService.AUDIT_CONFIG_FILE_NAME);
+ Path configPath = new Path(auditConfigPath);
+
+ assertTrue(HoodieCLI.fs.exists(configPath), "Config file should still exist");
+
+ String configContent;
+ try (InputStream inputStream = HoodieCLI.fs.open(configPath)) {
+ configContent = new String(FileIOUtils.readAsByteArray(inputStream));
+ }
+ JsonNode rootNode = OBJECT_MAPPER.readTree(configContent);
+ JsonNode enabledNode = rootNode.get(StorageLockProviderAuditService.STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD);
+
+ assertNotNull(enabledNode, "Config should contain enabled field");
+ assertFalse(enabledNode.asBoolean(), "Audit should be disabled");
+ }
+
+ /**
+ * Test disabling audit when no config file exists.
+ */
+ @Test
+ public void testDisableLockAuditNoConfigFile() {
+ Object result = shell.evaluate(() -> "locks audit disable");
+
+ assertAll("Disable command handles missing config file",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertEquals("Lock audit is already disabled (no configuration file found).", result.toString()));
+ }
+
+ /**
+ * Test disabling audit with keepAuditFiles=false option.
+ */
+ @Test
+ public void testDisableLockAuditWithoutKeepingFiles() throws Exception {
+ // First enable audit
+ shell.evaluate(() -> "locks audit enable");
+
+ // Create some audit files to be cleaned up
+ List<TransactionScenario> scenarios = new ArrayList<>();
+ List<LockAuditingCommand.AuditRecord> records = new ArrayList<>();
+ records.add(createAuditRecord("owner1", 1000L, 1100L, "START", 61000L));
+ records.add(createAuditRecord("owner1", 1000L, 1200L, "END", 61000L));
+ scenarios.add(new TransactionScenario("1000_owner1.jsonl", records));
+ createAuditFiles(scenarios);
+
+ // Disable with keepAuditFiles=false
+ Object result = shell.evaluate(() -> "locks audit disable --keepAuditFiles false");
+
+ assertAll("Disable command with keepAuditFiles=false",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Lock audit disabled successfully")),
+ () -> assertTrue(result.toString().contains("Audit files cleaned up at:")
+ || result.toString().contains("Some audit files may not have been cleaned up")));
+ }
+
+ /**
+ * Test that disable with keepAuditFiles=false actually cleans up files.
+ */
+ @Test
+ public void testDisableLockAuditCleansUpFiles() throws Exception {
+ // First enable audit
+ shell.evaluate(() -> "locks audit enable");
+
+ // Create some audit files
+ List<TransactionScenario> scenarios = new ArrayList<>();
+ List<LockAuditingCommand.AuditRecord> records = new ArrayList<>();
+ records.add(createAuditRecord("owner1", 1000L, 1100L, "START", 61000L));
+ records.add(createAuditRecord("owner1", 1000L, 1200L, "END", 61000L));
+ scenarios.add(new TransactionScenario("1000_owner1.jsonl", records));
+ createAuditFiles(scenarios);
+
+ // Verify files exist before disable
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditFolder = new Path(auditFolderPath);
+ FileStatus[] filesBefore = HoodieCLI.fs.listStatus(auditFolder);
+ long jsonlFilesBefore = Arrays.stream(filesBefore)
+ .filter(pathInfo -> pathInfo.isFile() && pathInfo.getPath().getName().endsWith(".jsonl"))
+ .count();
+ assertTrue(jsonlFilesBefore > 0, "Should have audit files before disable");
+
+ // Disable with keepAuditFiles=false - this should trigger cleanup
+ Object result = shell.evaluate(() -> "locks audit disable --keepAuditFiles false");
+
+ assertAll("Cleanup triggered by disable",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Lock audit disabled successfully")));
+ }
+
+ /**
+ * Test disable with keepAuditFiles=false when no audit files exist.
+ */
+ @Test
+ public void testDisableLockAuditWithoutKeepingFilesNoFiles() {
+ // First enable audit but don't create any audit files
+ shell.evaluate(() -> "locks audit enable");
+
+ // Disable with keepAuditFiles=false when no files exist
+ Object result = shell.evaluate(() -> "locks audit disable --keepAuditFiles false");
+
+ assertAll("Disable with cleanup when no files exist",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Lock audit disabled successfully")),
+ // Should handle the case where no files exist to clean up
+ () -> assertTrue(result.toString().contains("Audit files cleaned up at:")
+ || result.toString().contains("No audit files")
+ || result.toString().contains("nothing to cleanup")));
+ }
+
+ /**
+ * Test enabling audit multiple times (should overwrite).
+ */
+ @Test
+ public void testEnableLockAuditMultipleTimes() throws Exception {
+ // Enable first time
+ Object result1 = shell.evaluate(() -> "locks audit enable");
+ assertTrue(ShellEvaluationResultUtil.isSuccess(result1));
+ assertTrue(result1.toString().contains("Lock audit enabled successfully"));
+
+ // Enable second time (should succeed and overwrite)
+ Object result2 = shell.evaluate(() -> "locks audit enable");
+ assertTrue(ShellEvaluationResultUtil.isSuccess(result2));
+ assertTrue(result2.toString().contains("Lock audit enabled successfully"));
+
+ // Verify config is still correct
+ String lockFolderPath = String.format("%s%s.hoodie%s.locks", HoodieCLI.basePath, Path.SEPARATOR, Path.SEPARATOR);
+ String auditConfigPath = String.format("%s%s%s", lockFolderPath, Path.SEPARATOR, StorageLockProviderAuditService.AUDIT_CONFIG_FILE_NAME);
+ Path configPath = new Path(auditConfigPath);
+
+ String configContent;
+ try (InputStream inputStream = HoodieCLI.fs.open(configPath)) {
+ configContent = new String(FileIOUtils.readAsByteArray(inputStream));
+ }
+ JsonNode rootNode = OBJECT_MAPPER.readTree(configContent);
+ JsonNode enabledNode = rootNode.get(StorageLockProviderAuditService.STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD);
+
+ assertTrue(enabledNode.asBoolean(), "Audit should still be enabled");
+ }
+
+ // ==================== Validation Tests ====================
+
+ /**
+ * Test validation when no audit folder exists.
+ */
+ @Test
+ public void testValidateAuditLocksNoAuditFolder() {
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Validation handles missing audit folder",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: PASSED")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 0")),
+ () -> assertTrue(result.toString().contains("Issues Found: 0")),
+ () -> assertTrue(result.toString().contains("No audit folder found")));
+ }
+
+ /**
+ * Test validation when audit folder exists but no audit files.
+ */
+ @Test
+ public void testValidateAuditLocksNoAuditFiles() throws IOException {
+ // Create audit folder but no files
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditDir = new Path(auditFolderPath);
+ HoodieCLI.fs.mkdirs(auditDir);
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Validation handles no audit files",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: PASSED")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 0")),
+ () -> assertTrue(result.toString().contains("Issues Found: 0")),
+ () -> assertTrue(result.toString().contains("No audit files found")));
+ }
+
+ /**
+ * Test validation - No Issues (PASSED)
+ */
+ @Test
+ public void testValidateAuditLocksNoIssues() throws IOException {
+ long baseTime = System.currentTimeMillis();
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ // Transaction 1: Complete transaction
+ List<LockAuditingCommand.AuditRecord> records1 = new ArrayList<>();
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 100, "START", baseTime + 60000));
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 200, "RENEW", baseTime + 60000));
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 300, "END", baseTime + 60000));
+ scenarios.add(new TransactionScenario(baseTime + "_owner1.jsonl", records1));
+
+ // Transaction 2: Complete transaction starting after first one ends
+ List<LockAuditingCommand.AuditRecord> records2 = new ArrayList<>();
+ records2.add(createAuditRecord("owner2", baseTime + 500, baseTime + 600, "START", baseTime + 60000));
+ records2.add(createAuditRecord("owner2", baseTime + 500, baseTime + 700, "END", baseTime + 60000));
+ scenarios.add(new TransactionScenario((baseTime + 500) + "_owner2.jsonl", records2));
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Validation passes with no issues",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: PASSED")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 2")),
+ () -> assertTrue(result.toString().contains("Issues Found: 0")),
+ () -> assertTrue(result.toString().contains("successfully")));
+ }
+
+ /**
+ * Test validation - Single Unclosed Transaction (WARNING)
+ */
+ @Test
+ public void testValidateAuditLocksSingleUnclosedTransaction() throws IOException {
+ long baseTime = 1000000L;
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ // Single audit file without END record
+ List<LockAuditingCommand.AuditRecord> records1 = new ArrayList<>();
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 100, "START", baseTime + 200));
+ scenarios.add(new TransactionScenario(baseTime + "_owner1.jsonl", records1));
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Single unclosed transaction shows warning",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: WARNING")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 1")),
+ () -> assertTrue(result.toString().contains("Issues Found: 1")));
+ }
+
+ /**
+ * Test validation - Unclosed Transactions (WARNING)
+ */
+ @Test
+ public void testValidateAuditLocksUnclosedTransactions() throws IOException {
+ long baseTime = 1000000L;
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ // Transaction 1: Unclosed (effective end at expiration = baseTime + 200)
+ List<LockAuditingCommand.AuditRecord> records1 = new ArrayList<>();
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 100, "START", baseTime + 200));
+ scenarios.add(new TransactionScenario(baseTime + "_owner1.jsonl", records1));
+
+ // Transaction 2: Complete, starts after owner1's expiration
+ List<LockAuditingCommand.AuditRecord> records2 = new ArrayList<>();
+ records2.add(createAuditRecord("owner2", baseTime + 300, baseTime + 400, "START", baseTime + 60000));
+ records2.add(createAuditRecord("owner2", baseTime + 300, baseTime + 500, "END", baseTime + 60000));
+ scenarios.add(new TransactionScenario((baseTime + 300) + "_owner2.jsonl", records2));
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Unclosed transactions show warning",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: WARNING")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 2")),
+ () -> assertTrue(result.toString().contains("Issues Found: 1")),
+ () -> assertTrue(result.toString().contains("[WARNING]")),
+ () -> assertTrue(result.toString().contains("owner1.jsonl")),
+ () -> assertTrue(result.toString().contains("did not end gracefully")));
+ }
+
+ /**
+ * Test validation - Overlapping Transactions (FAILED)
+ */
+ @Test
+ public void testValidateAuditLocksOverlappingTransactions() throws IOException {
+ long baseTime = System.currentTimeMillis();
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ // Transaction 1: Ends after owner2 starts (overlapping)
+ List<LockAuditingCommand.AuditRecord> records1 = new ArrayList<>();
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 100, "START", baseTime + 60000));
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 500, "END", baseTime + 60000)); // Ends after owner2 starts
+ scenarios.add(new TransactionScenario(baseTime + "_owner1.jsonl", records1));
+
+ // Transaction 2: Starts before owner1 ends (overlapping)
+ List<LockAuditingCommand.AuditRecord> records2 = new ArrayList<>();
+ records2.add(createAuditRecord("owner2", baseTime + 200, baseTime + 300, "START", baseTime + 60000)); // Starts before owner1 ends
+ records2.add(createAuditRecord("owner2", baseTime + 200, baseTime + 400, "END", baseTime + 60000));
+ scenarios.add(new TransactionScenario((baseTime + 200) + "_owner2.jsonl", records2));
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Overlapping transactions show failure",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: FAILED")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 2")),
+ () -> assertTrue(result.toString().contains("Issues Found: 1")),
+ () -> assertTrue(result.toString().contains("[ERROR]")),
+ () -> assertTrue(result.toString().contains("owner1.jsonl")),
+ () -> assertTrue(result.toString().contains("overlaps with")),
+ () -> assertTrue(result.toString().contains("owner2.jsonl")));
+ }
+
+ /**
+ * Test validation - Mixed Issues (FAILED)
+ */
+ @Test
+ public void testValidateAuditLocksMixedIssues() throws IOException {
+ long baseTime = System.currentTimeMillis();
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ // Transaction 1: Unclosed transaction
+ List<LockAuditingCommand.AuditRecord> records1 = new ArrayList<>();
+ records1.add(createAuditRecord("owner1", baseTime, baseTime + 100, "START", baseTime + 60000));
+ // No END - unclosed
+ scenarios.add(new TransactionScenario(baseTime + "_owner1.jsonl", records1));
+
+ // Transaction 2: Overlaps with owner1
+ List<LockAuditingCommand.AuditRecord> records2 = new ArrayList<>();
+ records2.add(createAuditRecord("owner2", baseTime + 50, baseTime + 150, "START", baseTime + 60000)); // Overlaps with owner1
+ records2.add(createAuditRecord("owner2", baseTime + 50, baseTime + 250, "END", baseTime + 60000));
+ scenarios.add(new TransactionScenario((baseTime + 50) + "_owner2.jsonl", records2));
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Mixed issues show failure",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: FAILED")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 2")),
+ () -> assertTrue(result.toString().contains("Issues Found: 2")),
+ () -> assertTrue(result.toString().contains("[ERROR]")),
+ () -> assertTrue(result.toString().contains("[WARNING]")),
+ () -> assertTrue(result.toString().contains("overlaps with")),
+ () -> assertTrue(result.toString().contains("did not end gracefully")));
+ }
+
+ /**
+ * Test validation - Out of Order Filenames but Valid Transactions (PASSED)
+ */
+ @Test
+ public void testValidateAuditLocksOutOfOrderFilenames() throws IOException {
+ long baseTime = System.currentTimeMillis();
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ // File with later timestamp in name but contains earlier transaction
+ List<LockAuditingCommand.AuditRecord> records1 = new ArrayList<>();
+ records1.add(createAuditRecord("owner2", baseTime + 100, baseTime + 200, "START", baseTime + 60000)); // Actually starts first
+ records1.add(createAuditRecord("owner2", baseTime + 100, baseTime + 300, "END", baseTime + 60000));
+ scenarios.add(new TransactionScenario((baseTime + 2000) + "_owner2.jsonl", records1)); // Filename suggests later time
+
+ // File with earlier timestamp in name but contains later transaction
+ List<LockAuditingCommand.AuditRecord> records2 = new ArrayList<>();
+ records2.add(createAuditRecord("owner1", baseTime + 500, baseTime + 600, "START", baseTime + 60000)); // Actually starts second
+ records2.add(createAuditRecord("owner1", baseTime + 500, baseTime + 700, "END", baseTime + 60000));
+ scenarios.add(new TransactionScenario(baseTime + "_owner1.jsonl", records2)); // Filename suggests earlier time
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Out of order filenames but valid transactions pass",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: PASSED")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 2")),
+ () -> assertTrue(result.toString().contains("Issues Found: 0")),
+ () -> assertTrue(result.toString().contains("successfully")));
+ }
+
+ /**
+ * Test validation when no table is loaded.
+ */
+ @Test
+ public void testValidateAuditLocksNoTable() {
+ // Clear the base path to simulate no table loaded
+ String originalBasePath = HoodieCLI.basePath;
+ HoodieCLI.basePath = null;
+
+ try {
+ Object result = shell.evaluate(() -> "locks audit validate");
+ assertAll("Validation handles no table loaded",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertEquals("No Hudi table loaded. Please connect to a table first.", result.toString()));
+ } finally {
+ HoodieCLI.basePath = originalBasePath;
+ }
+ }
+
+ /**
+ * Test validation with malformed audit files.
+ */
+ @Test
+ public void testValidateAuditLocksMalformedFiles() throws IOException {
+ // Create audit directory
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditDir = new Path(auditFolderPath);
+ HoodieCLI.fs.mkdirs(auditDir);
+
+ // Create a malformed audit file
+ Path filePath = new Path(auditDir, "malformed_file.jsonl");
+ try (OutputStream outputStream = HoodieCLI.fs.create(filePath, true)) {
+ outputStream.write("invalid json content".getBytes());
+ }
+
+ Object result = shell.evaluate(() -> "locks audit validate");
+
+ assertAll("Validation handles malformed files",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Validation Result: FAILED")),
+ () -> assertTrue(result.toString().contains("Transactions Validated: 0")),
+ () -> assertTrue(result.toString().contains("Failed to parse any audit files")));
+ }
+
+ // ==================== Cleanup Tests ====================
+
+ /**
+ * Test cleanup when no audit folder exists.
+ */
+ @Test
+ public void testCleanupAuditLocksNoAuditFolder() {
+ Object result = shell.evaluate(() -> "locks audit cleanup");
+
+ assertAll("Cleanup handles missing audit folder",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertEquals("No audit folder found - nothing to cleanup.", result.toString()));
+ }
+
+ /**
+ * Test cleanup when audit folder exists but no audit files.
+ */
+ @Test
+ public void testCleanupAuditLocksNoAuditFiles() throws IOException {
+ // Create audit folder but no files
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditDir = new Path(auditFolderPath);
+ HoodieCLI.fs.mkdirs(auditDir);
+
+ Object result = shell.evaluate(() -> "locks audit cleanup");
+
+ assertAll("Cleanup handles no audit files",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("No audit files older than 7 days found")));
+ }
+
+ /**
+ * Test cleanup with dry run - test basic functionality.
+ */
+ @Test
+ public void testCleanupAuditLocksDryRun() throws IOException {
+ // Create some audit files (they will have current modification time)
+ long oldTime = System.currentTimeMillis() - (10 * 24 * 60 * 60 * 1000L); // 10 days ago
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ List<LockAuditingCommand.AuditRecord> records = new ArrayList<>();
+ records.add(createAuditRecord("owner1", oldTime, oldTime + 100, "START", oldTime + 60000));
+ records.add(createAuditRecord("owner1", oldTime, oldTime + 200, "END", oldTime + 60000));
+ scenarios.add(new TransactionScenario(oldTime + "_owner1.jsonl", records));
+
+ createAuditFiles(scenarios);
+
+ // Test default cleanup with dry run
+ Object result = shell.evaluate(() -> "locks audit cleanup --dryRun true");
+
+ assertAll("Dry run executes successfully",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ // Since files are created with current time, they should be too recent to delete
+ () -> assertTrue(result.toString().contains("No audit files older than")
+ || result.toString().contains("Dry run: Would delete")));
+
+ // Verify files still exist
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditFolder = new Path(auditFolderPath);
+ FileStatus[] filesAfter = HoodieCLI.fs.listStatus(auditFolder);
+ long jsonlFiles = Arrays.stream(filesAfter)
+ .filter(pathInfo -> pathInfo.isFile() && pathInfo.getPath().getName().endsWith(".jsonl"))
+ .count();
+ assertEquals(1, jsonlFiles, "Files should still exist after dry run");
+ }
+
+ /**
+ * Test cleanup with custom age threshold.
+ */
+ @Test
+ public void testCleanupAuditLocksCustomAge() throws IOException {
+ // Create some audit files
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ List<LockAuditingCommand.AuditRecord> records = new ArrayList<>();
+ records.add(createAuditRecord("owner1", 1000L, 1100L, "START", 61000L));
+ records.add(createAuditRecord("owner1", 1000L, 1200L, "END", 61000L));
+ scenarios.add(new TransactionScenario("1000_owner1.jsonl", records));
+
+ createAuditFiles(scenarios);
+
+ // Test with custom age threshold - files will be recent so won't be deleted
+ Object result = shell.evaluate(() -> "locks audit cleanup --ageDays 30 --dryRun true");
+
+ assertAll("Custom age threshold works",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("No audit files older than 30 days found")));
+ }
+
+ /**
+ * Test actual cleanup (not dry run).
+ */
+ @Test
+ public void testCleanupAuditLocksActualDelete() throws IOException {
+ // Create some audit files
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ List<LockAuditingCommand.AuditRecord> records = new ArrayList<>();
+ records.add(createAuditRecord("owner1", 1000L, 1100L, "START", 61000L));
+ records.add(createAuditRecord("owner1", 1000L, 1200L, "END", 61000L));
+ scenarios.add(new TransactionScenario("1000_owner1.jsonl", records));
+
+ createAuditFiles(scenarios);
+
+ // Test actual cleanup - files will be recent so shouldn't be deleted
+ Object result = shell.evaluate(() -> "locks audit cleanup --dryRun false");
+
+ assertAll("Actual cleanup executes successfully",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ // Recent files shouldn't be deleted
+ () -> assertTrue(result.toString().contains("No audit files older than 7 days found")));
+ }
+
+ /**
+ * Test cleanup with recent files (should not delete).
+ */
+ @Test
+ public void testCleanupAuditLocksRecentFiles() throws IOException {
+ // Create recent audit files
+ long recentTime = System.currentTimeMillis() - (2 * 24 * 60 * 60 * 1000L); // 2 days ago
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ List<LockAuditingCommand.AuditRecord> records = new ArrayList<>();
+ records.add(createAuditRecord("owner1", recentTime, recentTime + 100, "START", recentTime + 60000));
+ records.add(createAuditRecord("owner1", recentTime, recentTime + 200, "END", recentTime + 60000));
+ scenarios.add(new TransactionScenario(recentTime + "_owner1.jsonl", records));
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit cleanup");
+
+ assertAll("Recent files are not deleted",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("No audit files older than 7 days found")));
+ }
+
+ /**
+ * Test cleanup parameter validation - invalid age days.
+ */
+ @Test
+ public void testCleanupAuditLocksInvalidAgeDays() {
+ // Test with negative ageDays (Spring Shell treats this as invalid integer format)
+ Object resultNegative = shell.evaluate(() -> "locks audit cleanup --ageDays -1");
+
+ assertAll("Negative ageDays should be rejected",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(resultNegative)),
+ () -> assertNotNull(resultNegative.toString()),
+ () -> assertEquals("Error: ageDays must be a value greater than 0.", resultNegative.toString()));
+
+ // Test with invalid string to verify our parsing validation
+ Object resultInvalid = shell.evaluate(() -> "locks audit cleanup --ageDays abc");
+
+ assertAll("Invalid string ageDays should be rejected",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(resultInvalid)),
+ () -> assertNotNull(resultInvalid.toString()),
+ () -> assertEquals("Error: ageDays must be a value greater than 0.", resultInvalid.toString()));
+ }
+
+ /**
+ * Test cleanup when no table is loaded.
+ */
+ @Test
+ public void testCleanupAuditLocksNoTable() {
+ // Clear the base path to simulate no table loaded
+ String originalBasePath = HoodieCLI.basePath;
+ HoodieCLI.basePath = null;
+
+ Object result = shell.evaluate(() -> "locks audit cleanup");
+ assertAll("Cleanup handles no table loaded",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertEquals("No Hudi table loaded. Please connect to a table first.", result.toString()));
+ HoodieCLI.basePath = originalBasePath;
+ }
+
+ /**
+ * Test cleanup with multiple files.
+ */
+ @Test
+ public void testCleanupAuditLocksMixedFiles() throws IOException {
+ // Create multiple audit files
+ List<TransactionScenario> scenarios = new ArrayList<>();
+
+ // File 1
+ List<LockAuditingCommand.AuditRecord> records1 = new ArrayList<>();
+ records1.add(createAuditRecord("owner1", 1000L, 1100L, "START", 61000L));
+ records1.add(createAuditRecord("owner1", 1000L, 1200L, "END", 61000L));
+ scenarios.add(new TransactionScenario("1000_owner1.jsonl", records1));
+
+ // File 2
+ List<LockAuditingCommand.AuditRecord> records2 = new ArrayList<>();
+ records2.add(createAuditRecord("owner2", 2000L, 2100L, "START", 62000L));
+ records2.add(createAuditRecord("owner2", 2000L, 2200L, "END", 62000L));
+ scenarios.add(new TransactionScenario("2000_owner2.jsonl", records2));
+
+ createAuditFiles(scenarios);
+
+ Object result = shell.evaluate(() -> "locks audit cleanup --dryRun true");
+ System.out.println("DEBUG - Mixed files result: " + result.toString());
+
+ assertAll("Multiple files handled correctly",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("No audit files older than")
+ || result.toString().contains("cleanup")
+ || result.toString().contains("found")));
+ }
+
+ /**
+ * Test that disable method can call performAuditCleanup with ageDays=0 to delete all files.
+ * This validates that the internal validation allows ageDays >= 0 (not just > 0).
+ */
+ @Test
+ public void testDisableLockAuditWithAgeDaysZero() throws Exception {
+ // First enable audit
+ shell.evaluate(() -> "locks audit enable");
+
+ // Create some audit files to be cleaned up
+ List<TransactionScenario> scenarios = new ArrayList<>();
+ List<LockAuditingCommand.AuditRecord> records = new ArrayList<>();
+ records.add(createAuditRecord("owner1", 1000L, 1100L, "START", 61000L));
+ records.add(createAuditRecord("owner1", 1000L, 1200L, "END", 61000L));
+ scenarios.add(new TransactionScenario("1000_owner1.jsonl", records));
+ createAuditFiles(scenarios);
+
+ // Verify files exist before disable
+ String auditFolderPath = StorageLockProviderAuditService.getAuditFolderPath(HoodieCLI.basePath);
+ Path auditFolder = new Path(auditFolderPath);
+ FileStatus[] filesBefore = HoodieCLI.fs.listStatus(auditFolder);
+ long jsonlFilesBefore = Arrays.stream(filesBefore)
+ .filter(pathInfo -> pathInfo.isFile() && pathInfo.getPath().getName().endsWith(".jsonl"))
+ .count();
+ assertTrue(jsonlFilesBefore > 0, "Should have audit files before disable");
+
+ // Disable with keepAuditFiles=false, which internally calls performAuditCleanup(false, 0)
+ Object result = shell.evaluate(() -> "locks audit disable --keepAuditFiles false");
+
+ assertAll("Disable with ageDays=0 deletes all files",
+ () -> assertTrue(ShellEvaluationResultUtil.isSuccess(result)),
+ () -> assertNotNull(result.toString()),
+ () -> assertTrue(result.toString().contains("Lock audit disabled successfully")),
+ () -> assertTrue(result.toString().contains("cleaned up") || result.toString().contains("No audit files to clean up")));
+
+ // Verify files were cleaned up (ageDays=0 should delete all files)
+ FileStatus[] filesAfter = HoodieCLI.fs.listStatus(auditFolder);
+ long jsonlFilesAfter = Arrays.stream(filesAfter)
+ .filter(pathInfo -> pathInfo.isFile() && pathInfo.getPath().getName().endsWith(".jsonl"))
+ .count();
+ assertEquals(0, jsonlFilesAfter, "All audit files should be deleted when ageDays=0");
+ }
+}
\ No newline at end of file
diff --git a/hudi-client/hudi-client-common/pom.xml b/hudi-client/hudi-client-common/pom.xml
index 7a10ae12c35d..171c3384f82d 100644
--- a/hudi-client/hudi-client-common/pom.xml
+++ b/hudi-client/hudi-client-common/pom.xml
@@ -43,12 +43,6 @@
<artifactId>hudi-common</artifactId>
<version>${project.version}</version>
</dependency>
- <dependency>
- <groupId>org.apache.hudi</groupId>
- <artifactId>hudi-aws</artifactId>
- <version>${project.version}</version>
- <scope>provided</scope>
- </dependency>
<dependency>
<groupId>org.apache.hudi</groupId>
<artifactId>hudi-timeline-service</artifactId>
@@ -147,6 +141,12 @@
<type>test-jar</type>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.hudi</groupId>
+ <artifactId>hudi-hadoop-mr</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.apache.hudi</groupId>
<artifactId>hudi-tests-common</artifactId>
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/FileSystemBasedLockProvider.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/FileSystemBasedLockProvider.java
index da7e71a20580..4c883d753657 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/FileSystemBasedLockProvider.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/FileSystemBasedLockProvider.java
@@ -19,6 +19,7 @@
package org.apache.hudi.client.transaction.lock;
+import org.apache.hudi.client.transaction.lock.metrics.HoodieLockMetrics;
import org.apache.hudi.common.config.HoodieCommonConfig;
import org.apache.hudi.common.config.LockConfiguration;
import org.apache.hudi.common.config.TypedProperties;
@@ -90,6 +91,10 @@ public class FileSystemBasedLockProvider implements LockProvider<String>, Serial
}
}
+ public FileSystemBasedLockProvider(final LockConfiguration lockConfiguration, final Configuration configuration, final HoodieLockMetrics lockMetrics) {
+ this(lockConfiguration, configuration);
+ }
+
@Override
public void close() {
synchronized (LOCK_FILE_NAME) {
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/LockManager.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/LockManager.java
index 598f7cd70721..4a7b858465ed 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/LockManager.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/LockManager.java
@@ -26,8 +26,10 @@ import org.apache.hudi.common.lock.LockProvider;
import org.apache.hudi.common.util.ReflectionUtils;
import org.apache.hudi.config.HoodieLockConfig;
import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.exception.HoodieException;
import org.apache.hudi.exception.HoodieLockException;
+import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -107,7 +109,12 @@ public class LockManager implements Serializable, AutoCloseable {
*/
public void unlock() {
getLockProvider().unlock();
- metrics.updateLockHeldTimerMetrics();
+ try {
+ metrics.updateLockHeldTimerMetrics();
+ } catch (HoodieException e) {
+ LOG.error(String.format("Exception encountered when updating lock metrics: %s", e));
+ }
+ metrics.updateLockReleaseSuccessMetric();
close();
}
@@ -115,8 +122,19 @@ public class LockManager implements Serializable, AutoCloseable {
// Perform lazy initialization of lock provider only if needed
if (lockProvider == null) {
LOG.info("LockProvider " + writeConfig.getLockProviderClass());
- lockProvider = (LockProvider) ReflectionUtils.loadClass(writeConfig.getLockProviderClass(),
- lockConfiguration, hadoopConf.get());
+
+ // Try to load lock provider with HoodieLockMetrics constructor first
+ Class<?>[] metricsConstructorTypes = {LockConfiguration.class, Configuration.class, HoodieLockMetrics.class};
+ if (ReflectionUtils.hasConstructor(writeConfig.getLockProviderClass(), metricsConstructorTypes)) {
+ lockProvider = (LockProvider) ReflectionUtils.loadClass(writeConfig.getLockProviderClass(),
+ metricsConstructorTypes, lockConfiguration, hadoopConf.get(), metrics);
+ } else {
+ LOG.debug("LockProvider does not support HoodieLockMetrics param in constructor, falling back to standard constructor");
+ // Fallback to original constructor without metrics
+ lockProvider = (LockProvider) ReflectionUtils.loadClass(writeConfig.getLockProviderClass(),
+ new Class<?>[] {LockConfiguration.class, Configuration.class},
+ lockConfiguration, hadoopConf.get());
+ }
}
return lockProvider;
}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/StorageBasedLockProvider.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/StorageBasedLockProvider.java
new file mode 100644
index 000000000000..62912513ab9c
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/StorageBasedLockProvider.java
@@ -0,0 +1,651 @@
+/*
+ * 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.hudi.client.transaction.lock;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hudi.client.transaction.lock.audit.AuditOperationState;
+import org.apache.hudi.client.transaction.lock.audit.AuditService;
+import org.apache.hudi.client.transaction.lock.audit.AuditServiceFactory;
+import org.apache.hudi.client.transaction.lock.metrics.HoodieLockMetrics;
+import org.apache.hudi.client.transaction.lock.models.HeartbeatManager;
+import org.apache.hudi.client.transaction.lock.models.LockGetResult;
+import org.apache.hudi.client.transaction.lock.models.LockProviderHeartbeatManager;
+import org.apache.hudi.client.transaction.lock.models.LockUpsertResult;
+import org.apache.hudi.client.transaction.lock.models.StorageLockData;
+import org.apache.hudi.client.transaction.lock.models.StorageLockFile;
+import org.apache.hudi.common.config.LockConfiguration;
+import org.apache.hudi.common.config.TypedProperties;
+import org.apache.hudi.common.lock.LockProvider;
+import org.apache.hudi.common.lock.LockState;
+import org.apache.hudi.common.util.Functions;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.ReflectionUtils;
+import org.apache.hudi.common.util.VisibleForTesting;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.config.StorageBasedLockConfig;
+import org.apache.hudi.exception.HoodieLockException;
+import org.apache.hudi.exception.HoodieNotSupportedException;
+import org.apache.hudi.common.fs.StorageSchemes;
+
+import org.apache.hadoop.fs.Path;
+
+import org.jetbrains.annotations.NotNull;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.net.URI;
+import java.util.Objects;
+import java.util.Properties;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import static org.apache.hudi.common.config.LockConfiguration.DEFAULT_LOCK_ACQUIRE_RETRY_WAIT_TIME_IN_MILLIS;
+import static org.apache.hudi.common.lock.LockState.ACQUIRED;
+import static org.apache.hudi.common.lock.LockState.ACQUIRING;
+import static org.apache.hudi.common.lock.LockState.FAILED_TO_ACQUIRE;
+import static org.apache.hudi.common.lock.LockState.FAILED_TO_RELEASE;
+import static org.apache.hudi.common.lock.LockState.RELEASED;
+import static org.apache.hudi.common.lock.LockState.RELEASING;
+
+/**
+ * A distributed filesystem storage based lock provider. This {@link LockProvider} implementation
+ * leverages conditional writes to ensure transactional consistency for multi-writer scenarios.
+ * The underlying storage client interface {@link StorageLockClient} is pluggable so it can be implemented for any
+ * filesystem which supports conditional writes.
+ */
+@ThreadSafe
+public class StorageBasedLockProvider implements LockProvider<StorageLockFile> {
+
+ public static final String DEFAULT_TABLE_LOCK_FILE_NAME = "table_lock.json";
+ // How long to wait before retrying lock acquisition in blocking calls.
+ // Maximum expected clock drift between two nodes.
+ // This is similar idea as SkewAdjustingTimeGenerator.
+ // In reality, within a single cloud provider all nodes share the same ntp
+ // server
+ // therefore we do not expect drift more than a few ms.
+ // However, since our lock leases are pretty long, we can use a high buffer.
+ private static final long CLOCK_DRIFT_BUFFER_MS = 500;
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(StorageBasedLockProvider.class);
+
+ // Use for testing
+ private final Logger logger;
+
+ // The lock service implementation which interacts with storage
+ private final StorageLockClient storageLockClient;
+
+ private final long lockValiditySecs;
+ private final String ownerId;
+ private final String lockFilePath;
+ private final HeartbeatManager heartbeatManager;
+ private final transient Thread shutdownThread;
+ private final Option<HoodieLockMetrics> hoodieLockMetrics;
+ private Option<AuditService> auditService;
+ private final String basePath;
+
+ @GuardedBy("this")
+ private StorageLockFile currentLockObj = null;
+ @GuardedBy("this")
+ private boolean isClosed = false;
+
+ private synchronized void setLock(StorageLockFile lockObj) {
+ if (lockObj != null && !Objects.equals(lockObj.getOwner(), this.ownerId)) {
+ throw new HoodieLockException("Owners do not match. Current lock owner: " + this.ownerId + " lock path: "
+ + this.lockFilePath + " owner: " + lockObj.getOwner());
+ }
+ this.currentLockObj = lockObj;
+ }
+
+ /**
+ * Default constructor for StorageBasedLockProvider, required by LockManager
+ * to instantiate it using reflection.
+ *
+ * @param lockConfiguration The lock configuration, should be transformable into
+ * StorageBasedLockConfig
+ * @param conf Storage config, ignored.
+ */
+ public StorageBasedLockProvider(final LockConfiguration lockConfiguration, final Configuration conf) {
+ this(
+ UUID.randomUUID().toString(),
+ lockConfiguration.getConfig(),
+ LockProviderHeartbeatManager::new,
+ getStorageLockClientClassName(),
+ LOGGER,
+ null);
+ }
+
+ /**
+ * Constructor for StorageBasedLockProvider with HoodieLockMetrics support.
+ * This constructor allows lock providers to access metrics for fine-grained metrics collection.
+ *
+ * @param lockConfiguration The lock configuration, should be transformable into
+ * StorageBasedLockConfig
+ * @param conf Storage config, ignored.
+ * @param metrics HoodieLockMetrics instance for metrics collection
+ */
+ public StorageBasedLockProvider(final LockConfiguration lockConfiguration, final Configuration conf, final HoodieLockMetrics metrics) {
+ this(
+ UUID.randomUUID().toString(),
+ lockConfiguration.getConfig(),
+ LockProviderHeartbeatManager::new,
+ getStorageLockClientClassName(),
+ LOGGER,
+ metrics);
+ }
+
+ private static Functions.Function3<String, String, TypedProperties, StorageLockClient> getStorageLockClientClassName() {
+ return (ownerId, lockFilePath, lockConfig) -> {
+ try {
+ return (StorageLockClient) ReflectionUtils.loadClass(
+ getLockServiceClassName(new URI(lockFilePath).getScheme()),
+ new Class<?>[] {String.class, String.class, Properties.class},
+ new Object[] {ownerId, lockFilePath, lockConfig});
+ } catch (Throwable e) {
+ throw new HoodieLockException("Failed to load and initialize StorageLock", e);
+ }
+ };
+ }
+
+ private static @NotNull String getLockServiceClassName(String scheme) {
+ Option<StorageSchemes> schemeOptional = StorageSchemes.getStorageLockImplementationIfExists(scheme);
+ if (schemeOptional.isPresent()) {
+ return schemeOptional.get().getStorageLockClass();
+ } else {
+ throw new HoodieNotSupportedException("No implementation of StorageLock supports this scheme: " + scheme);
+ }
+ }
+
+ @VisibleForTesting
+ StorageBasedLockProvider(
+ String ownerId,
+ TypedProperties properties,
+ Functions.Function3<String, Long, Supplier<Boolean>, HeartbeatManager> heartbeatManagerLoader,
+ Functions.Function3<String, String, TypedProperties, StorageLockClient> storageLockClientLoader,
+ Logger logger,
+ HoodieLockMetrics hoodieLockMetrics) {
+ StorageBasedLockConfig config = new StorageBasedLockConfig.Builder().fromProperties(properties).build();
+ long heartbeatPollSeconds = config.getRenewIntervalSecs();
+ this.lockValiditySecs = config.getValiditySeconds();
+ this.basePath = config.getHudiTableBasePath();
+ String lockFolderPath = StorageLockClient.getLockFolderPath(basePath);
+ this.lockFilePath = new Path(lockFolderPath, DEFAULT_TABLE_LOCK_FILE_NAME).toString();
+ this.heartbeatManager = heartbeatManagerLoader.apply(ownerId, TimeUnit.SECONDS.toMillis(heartbeatPollSeconds), this::renewLock);
+ this.storageLockClient = storageLockClientLoader.apply(ownerId, lockFilePath, properties);
+ this.ownerId = ownerId;
+ this.logger = logger;
+ this.hoodieLockMetrics = Option.ofNullable(hoodieLockMetrics);
+ this.auditService = Option.empty(); // Will be created lazily on first lock acquisition
+ shutdownThread = new Thread(() -> shutdown(true));
+ Runtime.getRuntime().addShutdownHook(shutdownThread);
+ logger.info("Instantiated new storage-based lock provider, owner: {}, lockfilePath: {}", ownerId, lockFilePath);
+ }
+
+ // -----------------------------------------
+ // BASE METHODS
+ // -----------------------------------------
+
+ @Override
+ public synchronized StorageLockFile getLock() {
+ return currentLockObj;
+ }
+
+ /**
+ * Attempts to acquire the lock within the given timeout.
+ */
+ @Override
+ public boolean tryLock(long time, TimeUnit unit) {
+ long deadlineNanos = System.nanoTime() + unit.toNanos(time);
+
+ while (System.nanoTime() < deadlineNanos) {
+ try {
+ logDebugLockState(ACQUIRING);
+ if (tryLock()) {
+ return true;
+ }
+ Thread.sleep(Long.parseLong(DEFAULT_LOCK_ACQUIRE_RETRY_WAIT_TIME_IN_MILLIS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockInterruptedMetric);
+ throw new HoodieLockException(generateLockStateMessage(LockState.FAILED_TO_ACQUIRE), e);
+ }
+ }
+
+ return false;
+ }
+
+ @Override
+ public synchronized void close() {
+ shutdown(false);
+ }
+
+ @VisibleForTesting
+ synchronized void shutdown(boolean fromShutdownHook) {
+ if (fromShutdownHook) {
+ // Try to expire the lock from the shutdown hook.
+ if (!isClosed && actuallyHoldsLock()) {
+ tryExpireCurrentLock(true);
+ }
+ // Do not execute any further actions, mark closed
+ this.isClosed = true;
+ return;
+ } else {
+ try {
+ tryRemoveShutdownHook();
+ } catch (IllegalStateException e) {
+ logger.warn("Owner {}: Failed to remove shutdown hook, JVM is already shutting down.", ownerId, e);
+ }
+ }
+ try {
+ this.unlock();
+ } catch (Exception e) {
+ logger.error("Owner {}: Failed to unlock current lock.", ownerId, e);
+ }
+ try {
+ this.storageLockClient.close();
+ } catch (Exception e) {
+ logger.error("Owner {}: Lock service failed to close.", ownerId, e);
+ }
+ try {
+ this.heartbeatManager.close();
+ } catch (Exception e) {
+ logger.error("Owner {}: Heartbeat manager failed to close.", ownerId, e);
+ }
+ try {
+ this.auditService.ifPresent(auditService -> {
+ try {
+ auditService.close();
+ } catch (Exception e) {
+ logger.error("Owner {}: Audit service failed to close.", ownerId, e);
+ }
+ });
+ this.auditService = Option.empty();
+ } catch (Exception e) {
+ logger.error("Owner {}: Failed to close audit service.", ownerId, e);
+ }
+
+ this.isClosed = true;
+ }
+
+ @VisibleForTesting
+ void tryRemoveShutdownHook() {
+ Runtime.getRuntime().removeShutdownHook(shutdownThread);
+ }
+
+ private synchronized boolean isLockStillValid(StorageLockFile lock) {
+ if (lock.isExpired()) {
+ return false;
+ } else if (isCurrentTimeCertainlyOlderThanDistributedTime(lock.getValidUntilMs())) {
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockDanglingMetric);
+ logWarnLockState(ACQUIRING, "Found a dangling expired lock with owner " + lock.getOwner());
+ return false;
+ }
+
+ return true;
+ }
+
+ /**
+ * Attempts a single pass to acquire the lock (non-blocking).
+ *
+ * @return true if lock acquired, false otherwise
+ */
+ @Override
+ public synchronized boolean tryLock() {
+ assertHeartbeatManagerExists();
+ assertUnclosed();
+ logDebugLockState(ACQUIRING);
+ if (actuallyHoldsLock()) {
+ // Supports reentrant locks
+ return true;
+ }
+
+ if (this.heartbeatManager.hasActiveHeartbeat()) {
+ logger.error("Detected broken invariant: there is an active heartbeat without a lock being held.");
+ // Breach of object invariant - we should never have an active heartbeat without
+ // holding a lock.
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockProviderFatalErrorMetric);
+ throw new HoodieLockException(generateLockStateMessage(FAILED_TO_ACQUIRE));
+ }
+
+ Pair<LockGetResult, Option<StorageLockFile>> latestLock = this.storageLockClient.readCurrentLockFile();
+ if (latestLock.getLeft() == LockGetResult.UNKNOWN_ERROR) {
+ logInfoLockState(FAILED_TO_ACQUIRE, "Failed to get the latest lock status");
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockStateUnknownMetric);
+ // We were not able to determine whether a lock was present.
+ return false;
+ }
+
+ if (latestLock.getLeft() == LockGetResult.SUCCESS && isLockStillValid(latestLock.getRight().get())) {
+ String msg = String.format("Lock already held by %s", latestLock.getRight().get().getOwner());
+ // Lock held by others.
+ logInfoLockState(FAILED_TO_ACQUIRE, msg);
+ return false;
+ }
+
+ // Try to acquire the lock
+ long acquisitionTimestamp = getCurrentEpochMs();
+ long lockExpirationMs = calculateLockExpiration(acquisitionTimestamp);
+ StorageLockData newLockData = new StorageLockData(false, lockExpirationMs, ownerId);
+ Pair<LockUpsertResult, Option<StorageLockFile>> lockUpdateStatus = this.storageLockClient.tryUpsertLockFile(
+ newLockData,
+ latestLock.getRight());
+ if (lockUpdateStatus.getLeft() != LockUpsertResult.SUCCESS) {
+ // failed to acquire the lock, indicates concurrent contention
+ logInfoLockState(FAILED_TO_ACQUIRE);
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockAcquirePreconditionFailureMetric);
+ return false;
+ }
+ this.setLock(lockUpdateStatus.getRight().get());
+ hoodieLockMetrics.ifPresent(metrics -> metrics.updateLockExpirationDeadlineMetric(
+ (int) (lockUpdateStatus.getRight().get().getValidUntilMs() - getCurrentEpochMs())));
+
+ // There is a remote chance that
+ // - after lock is acquired but before heartbeat starts the lock is expired.
+ // - lock is acquired and heartbeat is up yet it does not run timely before the
+ // lock is expired
+ // It is mitigated by setting the lock validity period to a reasonably long
+ // period to survive until heartbeat comes, plus
+ // set the heartbeat interval relatively small enough.
+ if (!this.heartbeatManager.startHeartbeatForThread(Thread.currentThread())) {
+ // Precondition "no active heartbeat" is checked previously, so when
+ // startHeartbeatForThread returns false,
+ // we are confident no heartbeat thread is running.
+ logErrorLockState(RELEASING, "We were unable to start the heartbeat!");
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockProviderFatalErrorMetric);
+ tryExpireCurrentLock(false);
+ return false;
+ }
+
+ logInfoLockState(ACQUIRED);
+
+ // Create audit service lazily on first successful lock acquisition if auditing is enabled
+ if (auditService.isEmpty()) {
+ auditService = AuditServiceFactory.createLockProviderAuditService(
+ ownerId, basePath, storageLockClient, acquisitionTimestamp,
+ this::calculateLockExpiration, this::actuallyHoldsLock);
+ }
+
+ recordAuditOperation(AuditOperationState.START, acquisitionTimestamp);
+ return true;
+ }
+
+ /**
+ * Determines whether this provider currently holds a valid lock.
+ *
+ * <p>
+ * This method checks both the existence of a lock object and its validity. A
+ * lock is considered
+ * valid only if it exists and has not expired according to its timestamp.
+ *
+ * @return {@code true} if this provider holds a valid lock, {@code false}
+ * otherwise
+ */
+ private boolean actuallyHoldsLock() {
+ return believesLockMightBeHeld() && isLockStillValid(getLock());
+ }
+
+ /**
+ * Checks if this provider has a non-null lock object reference.
+ *
+ * <p>
+ * A non-null lock object indicates that this provider has previously
+ * **successfully** acquired a lock via
+ * StorageBasedLockProvider##lock and has not yet **successfully** released
+ * it via StorageBasedLockProvider#unlock().
+ * It is merely an indicator that the lock might be held by this provider. To
+ * truly certify we are the owner of the lock,
+ * StorageBasedLockProvider#actuallyHoldsLock should be used.
+ *
+ * @return {@code true} if this provider has a non-null lock object,
+ * {@code false} otherwise
+ * @see StorageBasedLockProvider#actuallyHoldsLock()
+ */
+ private boolean believesLockMightBeHeld() {
+ return this.getLock() != null;
+ }
+
+ /**
+ * Unlock by marking our current lock file "expired": true.
+ */
+ @Override
+ public synchronized void unlock() {
+ assertHeartbeatManagerExists();
+ if (!believesLockMightBeHeld()) {
+ return;
+ }
+ boolean believesNoLongerHoldsLock = true;
+
+ // Try to stop the heartbeat first
+ if (heartbeatManager.hasActiveHeartbeat()) {
+ logger.debug("Owner {}: Gracefully shutting down heartbeat.", ownerId);
+ believesNoLongerHoldsLock &= heartbeatManager.stopHeartbeat(true);
+ }
+
+ // Then expire the current lock.
+ believesNoLongerHoldsLock &= tryExpireCurrentLock(false);
+ if (!believesNoLongerHoldsLock) {
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockReleaseFailureMetric);
+ throw new HoodieLockException(generateLockStateMessage(FAILED_TO_RELEASE));
+ }
+ }
+
+ private void assertHeartbeatManagerExists() {
+ if (heartbeatManager == null) {
+ // broken function precondition.
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockProviderFatalErrorMetric);
+ throw new HoodieLockException("Unexpected null heartbeatManager");
+ }
+ }
+
+ private void assertUnclosed() {
+ if (this.isClosed) {
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockProviderFatalErrorMetric);
+ throw new HoodieLockException("Lock provider already closed");
+ }
+ }
+
+ /**
+ * Tries to expire the currently held lock.
+ *
+ * @param fromShutdownHook Whether we are attempting best effort quick unlock from shutdown hook.
+ * @return True if we were successfully able to upload an expired lock.
+ */
+ private synchronized boolean tryExpireCurrentLock(boolean fromShutdownHook) {
+ // It does not make sense to have heartbeat alive extending the lock lease while
+ // here we are trying
+ // to expire the lock.
+ if (!fromShutdownHook && heartbeatManager.hasActiveHeartbeat()) {
+ // broken function precondition.
+ throw new HoodieLockException("Must stop heartbeat before expire lock file");
+ }
+ logDebugLockState(RELEASING);
+ // Upload metadata that will unlock this lock.
+ StorageLockData expiredLockData = new StorageLockData(true, this.getLock().getValidUntilMs(), ownerId);
+ Pair<LockUpsertResult, Option<StorageLockFile>> result;
+ long lockExpirationTimeMs = System.currentTimeMillis();
+ result = this.storageLockClient.tryUpsertLockFile(expiredLockData, Option.of(this.getLock()));
+ switch (result.getLeft()) {
+ case UNKNOWN_ERROR:
+ // Here we do not know the state of the lock.
+ logErrorLockState(FAILED_TO_RELEASE, "Lock state is unknown.");
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockStateUnknownMetric);
+ return false;
+ case SUCCESS:
+ logInfoLockState(RELEASED);
+ recordAuditOperation(AuditOperationState.END, lockExpirationTimeMs);
+ setLock(null);
+ return true;
+ case ACQUIRED_BY_OTHERS:
+ // Lock was acquired by others, indicating heartbeat failure during lock hold period.
+ logErrorLockState(FAILED_TO_RELEASE, "lock was acquired by others, indicating heartbeat failure.");
+ setLock(null);
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockAcquiredByOthersErrorMetric);
+ return false;
+ default:
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockReleaseFailureMetric);
+ throw new HoodieLockException("Unexpected lock update result: " + result.getLeft());
+ }
+ }
+
+ /**
+ * Renews (heartbeats) the current lock if we are the holder, it forcefully set
+ * the expiration flag
+ * to false and the lock expiration time to a later time in the future.
+ *
+ * @return True if we successfully renewed the lock, false if not.
+ */
+ @VisibleForTesting
+ protected synchronized boolean renewLock() {
+ try {
+ // If we don't hold the lock, no-op.
+ if (!believesLockMightBeHeld()) {
+ logger.warn("Owner {}: Cannot renew, no lock held by this process", ownerId);
+ // No need to extend lock lease.
+ return false;
+ }
+
+ long oldExpirationMs = getLock().getValidUntilMs();
+ // Attempt conditional update, extend lock. There are 3 cases:
+ // 1. Happy case: lock has not expired yet, we extend the lease to a longer
+ // period.
+ // 2. Corner case 1: lock is expired and is acquired by others, lock renewal
+ // failed with ACQUIRED_BY_OTHERS.
+ // 3. Corner case 2: lock is expired but no one has acquired it yet, lock
+ // renewal "revived" the expired lock.
+ // Please note we expect the corner cases almost never happens.
+ // Action taken for corner case 2 is just a best effort mitigation. At least it
+ // prevents further data corruption by
+ // letting someone else acquire the lock.
+ long acquisitionTimestamp = getCurrentEpochMs();
+ long lockExpirationMs = calculateLockExpiration(acquisitionTimestamp);
+ Pair<LockUpsertResult, Option<StorageLockFile>> currentLock = this.storageLockClient.tryUpsertLockFile(
+ new StorageLockData(false, lockExpirationMs, ownerId),
+ Option.of(getLock()));
+ switch (currentLock.getLeft()) {
+ case ACQUIRED_BY_OTHERS:
+ logger.error("Owner {}: Unable to renew lock as it is acquired by others.", ownerId);
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockAcquiredByOthersErrorMetric);
+ // No need to extend lock lease anymore.
+ return false;
+ case UNKNOWN_ERROR:
+ // This could be transient, but unclear, we will let the heartbeat continue
+ // normally.
+ // If the next heartbeat run identifies our lock has expired we will error out.
+ logger.warn("Owner {}: Unable to renew lock due to unknown error, could be transient.", ownerId);
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockStateUnknownMetric);
+ // Let heartbeat retry later.
+ return true;
+ case SUCCESS:
+ // Only positive outcome
+ this.setLock(currentLock.getRight().get());
+ hoodieLockMetrics.ifPresent(metrics -> metrics.updateLockExpirationDeadlineMetric(
+ (int) (oldExpirationMs - getCurrentEpochMs())));
+ logger.info("Owner {}: Lock renewal successful. The renewal completes {} ms before expiration for lock {}.",
+ ownerId, oldExpirationMs - getCurrentEpochMs(), lockFilePath);
+ recordAuditOperation(AuditOperationState.RENEW, acquisitionTimestamp);
+ // Let heartbeat continue to renew lock lease again later.
+ return true;
+ default:
+ throw new HoodieLockException("Unexpected lock update result: " + currentLock.getLeft());
+ }
+ } catch (Exception e) {
+ logger.error("Owner {}: Exception occurred while renewing lock", ownerId, e);
+ hoodieLockMetrics.ifPresent(HoodieLockMetrics::updateLockProviderFatalErrorMetric);
+ return false;
+ }
+ }
+
+ // ---------
+ // Utilities
+ // ---------
+
+ /**
+ * Method to calculate whether a timestamp from a distributed source has
+ * definitively occurred yet.
+ */
+ protected boolean isCurrentTimeCertainlyOlderThanDistributedTime(long epochMs) {
+ return getCurrentEpochMs() > epochMs + CLOCK_DRIFT_BUFFER_MS;
+ }
+
+ private String generateLockStateMessage(LockState state) {
+ String threadName = Thread.currentThread().getName();
+ return String.format(
+ "Owner %s: Lock file path %s, Thread %s, Storage based lock state %s",
+ ownerId,
+ lockFilePath,
+ threadName,
+ state.toString());
+ }
+
+ private static final String LOCK_STATE_LOGGER_MSG = "Owner {}: Lock file path {}, Thread {}, Storage based lock state {}";
+ private static final String LOCK_STATE_LOGGER_MSG_WITH_INFO = "Owner {}: Lock file path {}, Thread {}, Storage based lock state {}, {}";
+
+ private void logDebugLockState(LockState state) {
+ logger.debug(LOCK_STATE_LOGGER_MSG, ownerId, lockFilePath, Thread.currentThread(), state);
+ }
+
+ private void logInfoLockState(LockState state) {
+ logger.info(LOCK_STATE_LOGGER_MSG, ownerId, lockFilePath, Thread.currentThread(), state);
+ }
+
+ private void logInfoLockState(LockState state, String msg) {
+ logger.info(LOCK_STATE_LOGGER_MSG_WITH_INFO, ownerId, lockFilePath, Thread.currentThread(), state, msg);
+ }
+
+ private void logWarnLockState(LockState state, String msg) {
+ logger.warn(LOCK_STATE_LOGGER_MSG_WITH_INFO, ownerId, lockFilePath, Thread.currentThread(), state, msg);
+ }
+
+ private void logErrorLockState(LockState state, String msg) {
+ logger.error(LOCK_STATE_LOGGER_MSG_WITH_INFO, ownerId, lockFilePath, Thread.currentThread(), state, msg);
+ }
+
+ @VisibleForTesting
+ long getCurrentEpochMs() {
+ return System.currentTimeMillis();
+ }
+
+ /**
+ * Calculates the lock expiration time based on given timestamp and validity period.
+ * This is the shared function used by both the lock provider and audit service.
+ *
+ * @param timestamp The base timestamp to calculate expiration from
+ * @return Lock expiration time in milliseconds
+ */
+ private long calculateLockExpiration(long timestamp) {
+ return timestamp + TimeUnit.SECONDS.toMillis(lockValiditySecs);
+ }
+
+ /**
+ * Helper method to record audit operations.
+ */
+ private void recordAuditOperation(AuditOperationState state, long timestamp) {
+ auditService.ifPresent(service -> {
+ try {
+ service.recordOperation(state, timestamp);
+ } catch (Exception e) {
+ // Log but don't fail the lock operation due to recording failures
+ logger.warn("Owner {}: Failed to record audit operation {}: {}", ownerId, state, e.getMessage());
+ }
+ });
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/StorageLockClient.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/StorageLockClient.java
new file mode 100644
index 000000000000..67755f36c5d3
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/StorageLockClient.java
@@ -0,0 +1,122 @@
+/*
+ * 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.hudi.client.transaction.lock;
+
+import org.apache.hudi.client.transaction.lock.models.LockUpsertResult;
+import org.apache.hudi.client.transaction.lock.models.StorageLockData;
+import org.apache.hudi.client.transaction.lock.models.StorageLockFile;
+import org.apache.hudi.client.transaction.lock.models.LockGetResult;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.exception.HoodieLockException;
+
+import org.apache.hadoop.fs.Path;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+
+import static org.apache.hudi.common.table.HoodieTableMetaClient.LOCKS_FOLDER_NAME;
+
+/**
+ * Defines a contract for a service which should be able to perform conditional writes to object storage.
+ * It expects to be interacting with a single lock file per context (table), and will be competing with other instances
+ * to perform writes, so it should handle these cases accordingly (using conditional writes).
+ */
+public interface StorageLockClient extends AutoCloseable {
+ /**
+ * Tries to create or update a lock file.
+ * All non pre-condition failure related errors should be returned as UNKNOWN_ERROR.
+ *
+ * @param newLockData The new data to write to the lock file
+ * @param previousLockFile The previous lock file
+ * @return A pair containing the result state and the new lock file (if successful)
+ */
+ Pair<LockUpsertResult, Option<StorageLockFile>> tryUpsertLockFile(
+ StorageLockData newLockData,
+ Option<StorageLockFile> previousLockFile);
+
+ /**
+ * Reads the current lock file.
+ *
+ * @return The lock retrieve result and the current lock file if successfully retrieved.
+ */
+ Pair<LockGetResult, Option<StorageLockFile>> readCurrentLockFile();
+
+ /**
+ * Reads an object from the specified path.
+ * This method is intended for reading small files (e.g., audit config, feature flags, JSON objects)
+ * and loads the entire file into memory. Do not use for large files.
+ *
+ * @param filePath The path to the object file to read
+ * @param checkExistsFirst If true, performs a lightweight existence check before attempting to read.
+ * Recommended when the file is unlikely to exist (e.g., optional configs).
+ * @return An Option containing the content as a string if successful, Option.empty() otherwise
+ */
+ Option<String> readObject(String filePath, boolean checkExistsFirst);
+
+ /**
+ * Writes an object to the specified path.
+ * This method is intended for writing small files (e.g., audit logs, configuration files)
+ * and should not be used for large files.
+ *
+ * @param filePath The path where the object should be written
+ * @param content The content to write as a string
+ * @return true if the write was successful, false otherwise
+ */
+ boolean writeObject(String filePath, String content);
+
+ /**
+ * Gets the lock folder path for the given base path.
+ * This is a static utility method that can be used without creating an instance.
+ *
+ * @param basePath The base path of the Hudi table
+ * @return The lock folder path (e.g., "s3://bucket/table/.hoodie/.locks")
+ */
+ static String getLockFolderPath(String basePath) {
+ return new Path(basePath, LOCKS_FOLDER_NAME).toString();
+ }
+
+ /**
+ * Parses a URI and returns bucket name and path as a Pair.
+ * This is a shared utility method for all storage lock client implementations.
+ *
+ * @param uriString The URI string to parse
+ * @return A Pair containing bucket name (left) and path (right)
+ * @throws HoodieLockException if URI parsing fails or components are invalid
+ */
+ static Pair<String, String> parseBucketAndPath(String uriString) {
+ try {
+ URI uri = new URI(uriString);
+ String bucketName = uri.getAuthority();
+ String path = uri.getPath().replaceFirst("/", "");
+
+ if (StringUtils.isNullOrEmpty(bucketName)) {
+ throw new IllegalArgumentException("URI does not contain a valid bucket name.");
+ }
+ if (StringUtils.isNullOrEmpty(path)) {
+ throw new IllegalArgumentException("URI does not contain a valid path.");
+ }
+
+ return Pair.of(bucketName, path);
+ } catch (URISyntaxException e) {
+ throw new HoodieLockException("Failed to parse URI: " + uriString, e);
+ }
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/ZookeeperBasedImplicitBasePathLockProvider.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/ZookeeperBasedImplicitBasePathLockProvider.java
index b7c2da005f62..5d7c1125db7d 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/ZookeeperBasedImplicitBasePathLockProvider.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/ZookeeperBasedImplicitBasePathLockProvider.java
@@ -30,7 +30,7 @@ import org.slf4j.LoggerFactory;
import javax.annotation.concurrent.NotThreadSafe;
-import static org.apache.hudi.aws.utils.S3Utils.s3aToS3;
+import static org.apache.hudi.common.fs.FSUtils.s3aToS3;
/**
* A zookeeper based lock. This {@link LockProvider} implementation allows to lock table operations
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditOperationState.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditOperationState.java
new file mode 100644
index 000000000000..700a8dd79be7
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditOperationState.java
@@ -0,0 +1,39 @@
+/*
+ * 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.hudi.client.transaction.lock.audit;
+
+/**
+ * Enumeration of audit operation states for different operations.
+ */
+public enum AuditOperationState {
+ /**
+ * Operation started.
+ */
+ START,
+
+ /**
+ * Operation renewal/heartbeat.
+ */
+ RENEW,
+
+ /**
+ * Operation end.
+ */
+ END
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditService.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditService.java
new file mode 100644
index 000000000000..bdc062f519af
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditService.java
@@ -0,0 +1,41 @@
+/*
+ * 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.hudi.client.transaction.lock.audit;
+
+import org.apache.hudi.ApiMaturityLevel;
+import org.apache.hudi.PublicAPIClass;
+import org.apache.hudi.PublicAPIMethod;
+
+/**
+ * Generic audit service interface for tracking operation lifecycles.
+ * Provides a single method for recording all types of audit operations.
+ */
+@PublicAPIClass(maturity = ApiMaturityLevel.EVOLVING)
+public interface AuditService extends AutoCloseable {
+
+ /**
+ * Records an audit operation with the given state and timestamp.
+ *
+ * @param state The type of operation (START, RENEW, END)
+ * @param timestamp When the operation occurred
+ * @throws Exception if the operation cannot be recorded
+ */
+ @PublicAPIMethod(maturity = ApiMaturityLevel.EVOLVING)
+ void recordOperation(AuditOperationState state, long timestamp) throws Exception;
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditServiceFactory.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditServiceFactory.java
new file mode 100644
index 000000000000..871330886916
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/AuditServiceFactory.java
@@ -0,0 +1,119 @@
+/*
+ * 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.hudi.client.transaction.lock.audit;
+
+import org.apache.hudi.client.transaction.lock.StorageLockClient;
+import org.apache.hudi.common.util.Option;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * Generic factory for creating audit services.
+ * This factory determines whether auditing is enabled by checking configuration files.
+ */
+public class AuditServiceFactory {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AuditServiceFactory.class);
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ /**
+ * Creates a lock provider audit service instance by checking the audit configuration file.
+ * This method reads the audit configuration to determine if auditing is enabled.
+ * If auditing is enabled, it returns a StorageLockProviderAuditService instance.
+ *
+ * @param ownerId The owner ID for the lock provider
+ * @param basePath The base path of the Hudi table
+ * @param storageLockClient The storage lock client to use for reading configuration
+ * @param transactionStartTime The timestamp when the transaction started (lock acquired)
+ * @param lockExpirationFunction Function that takes a timestamp and returns the lock expiration time
+ * @param lockHeldSupplier Supplier that provides whether the lock is currently held
+ * @return An Option containing the audit service if enabled, Option.empty() otherwise
+ */
+ public static Option<AuditService> createLockProviderAuditService(
+ String ownerId,
+ String basePath,
+ StorageLockClient storageLockClient,
+ long transactionStartTime,
+ Function<Long, Long> lockExpirationFunction,
+ Supplier<Boolean> lockHeldSupplier) {
+
+ if (!isAuditEnabled(basePath, storageLockClient)) {
+ return Option.empty();
+ }
+
+ return Option.of(new StorageLockProviderAuditService(
+ basePath,
+ ownerId,
+ transactionStartTime,
+ storageLockClient,
+ lockExpirationFunction,
+ lockHeldSupplier));
+ }
+
+ /**
+ * Checks if audit is enabled by reading the audit configuration file.
+ * This method looks for the audit configuration file at the expected location
+ * and parses it to determine if auditing is enabled. Returns false if the
+ * configuration file doesn't exist or cannot be read.
+ *
+ * @param basePath The base path of the Hudi table
+ * @param storageLockClient The storage lock client to use for reading configuration
+ * @return true if audit is enabled, false otherwise
+ */
+ private static boolean isAuditEnabled(String basePath, StorageLockClient storageLockClient) {
+ try {
+ // Construct the audit config path using the utility method
+ String auditConfigPath = StorageLockProviderAuditService.getAuditConfigPath(basePath);
+
+ LOG.debug("Checking for audit configuration at: {}", auditConfigPath);
+
+ // Use the readObject method to read the config file
+ // Pass true for checkExistsFirst since audit config is rarely present (99% miss rate)
+ Option<String> jsonContent = storageLockClient.readObject(auditConfigPath, true);
+
+ if (jsonContent.isPresent()) {
+ LOG.debug("Audit configuration file found, parsing content");
+ JsonNode rootNode = OBJECT_MAPPER.readTree(jsonContent.get());
+ JsonNode enabledNode = rootNode.get(StorageLockProviderAuditService.STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD);
+
+ boolean isEnabled = enabledNode != null && enabledNode.asBoolean(false);
+
+ if (isEnabled) {
+ LOG.info("Audit logging is ENABLED for lock operations at base path: {}", basePath);
+ } else {
+ LOG.info("Audit configuration present but audit logging is DISABLED");
+ }
+
+ return isEnabled;
+ }
+
+ LOG.debug("No audit configuration file found at: {}", auditConfigPath);
+ return false;
+ } catch (Exception e) {
+ LOG.error("Error reading audit configuration from base path: {}. Audit will be disabled.", basePath, e);
+ return false;
+ }
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/StorageLockProviderAuditService.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/StorageLockProviderAuditService.java
new file mode 100644
index 000000000000..b7c09479648b
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/audit/StorageLockProviderAuditService.java
@@ -0,0 +1,177 @@
+/*
+ * 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.hudi.client.transaction.lock.audit;
+
+import org.apache.hudi.client.transaction.lock.StorageLockClient;
+
+import org.apache.hadoop.fs.Path;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * Storage-based audit service implementation for lock provider operations.
+ * Writes audit records to a single JSONL file per transaction to track lock lifecycle events.
+ */
+public class StorageLockProviderAuditService implements AuditService {
+
+ private static final Logger LOG = LoggerFactory.getLogger(StorageLockProviderAuditService.class);
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ // Audit configuration constants
+ public static final String AUDIT_FOLDER_NAME = "audit";
+ public static final String AUDIT_CONFIG_FILE_NAME = "audit_enabled.json";
+ public static final String STORAGE_LOCK_AUDIT_SERVICE_ENABLED_FIELD = "STORAGE_LOCK_AUDIT_SERVICE_ENABLED";
+
+ /**
+ * Constructs the full path to the audit configuration file for a given table.
+ *
+ * @param basePath The base path of the Hudi table
+ * @return The full path to the audit_enabled.json configuration file
+ */
+ public static String getAuditConfigPath(String basePath) {
+ String lockFolderPath = StorageLockClient.getLockFolderPath(basePath);
+ return new Path(lockFolderPath, AUDIT_CONFIG_FILE_NAME).toString();
+ }
+
+ /**
+ * Constructs the full path to the audit folder for a given table.
+ *
+ * @param basePath The base path of the Hudi table
+ * @return The full path to the audit folder where audit files are stored
+ */
+ public static String getAuditFolderPath(String basePath) {
+ String lockFolderPath = StorageLockClient.getLockFolderPath(basePath);
+ return new Path(lockFolderPath, AUDIT_FOLDER_NAME).toString();
+ }
+
+ private final String ownerId;
+ private final long transactionStartTime;
+ private final String auditFilePath;
+ private final StorageLockClient storageLockClient;
+ private final Function<Long, Long> lockExpirationFunction;
+ private final Supplier<Boolean> lockHeldSupplier;
+ private final StringBuilder auditBuffer;
+
+ /**
+ * Creates a new StorageLockProviderAuditService instance.
+ *
+ * @param basePath The base path where audit files will be written
+ * @param ownerId The full owner ID for the lock
+ * @param transactionStartTime The timestamp when the transaction started (lock acquired)
+ * @param storageLockClient The storage client for writing audit files
+ * @param lockExpirationFunction Function that takes a timestamp and returns the lock expiration time
+ * @param lockHeldSupplier Supplier that provides whether the lock is currently held
+ */
+ public StorageLockProviderAuditService(
+ String basePath,
+ String ownerId,
+ long transactionStartTime,
+ StorageLockClient storageLockClient,
+ Function<Long, Long> lockExpirationFunction,
+ Supplier<Boolean> lockHeldSupplier) {
+ this.ownerId = ownerId;
+ this.transactionStartTime = transactionStartTime;
+ this.storageLockClient = storageLockClient;
+ this.lockExpirationFunction = lockExpirationFunction;
+ this.lockHeldSupplier = lockHeldSupplier;
+ this.auditBuffer = new StringBuilder();
+
+ // Generate audit file path: <txn-start>_<full-owner-id>.jsonl
+ String filename = String.format("%d_%s.jsonl", transactionStartTime, ownerId);
+ this.auditFilePath = new Path(getAuditFolderPath(basePath), filename).toString();
+
+ LOG.debug("Initialized audit service for transaction starting at {} with file: {}",
+ transactionStartTime, auditFilePath);
+ }
+
+ /**
+ * Records an audit operation with the current timestamp and state.
+ *
+ * @param state The audit operation state to record
+ * @param timestamp The timestamp when the operation occurred
+ * @throws Exception if there's an error creating or writing the audit record
+ */
+ @Override
+ public synchronized void recordOperation(AuditOperationState state, long timestamp) throws Exception {
+ // Create audit record
+ Map<String, Object> auditRecord = new HashMap<>();
+ auditRecord.put("ownerId", ownerId);
+ auditRecord.put("transactionStartTime", transactionStartTime);
+ auditRecord.put("timestamp", timestamp);
+ auditRecord.put("state", state.name());
+ auditRecord.put("lockExpiration", lockExpirationFunction.apply(timestamp));
+ auditRecord.put("lockHeld", lockHeldSupplier.get());
+
+ // Convert to JSON and append newline for JSONL format
+ String jsonLine = OBJECT_MAPPER.writeValueAsString(auditRecord) + "\n";
+
+ // Append to buffer
+ auditBuffer.append(jsonLine);
+
+ // Write the accumulated audit records to file
+ writeAuditFile();
+
+ LOG.debug("Recorded audit operation: state={}, timestamp={}, file={}",
+ state, timestamp, auditFilePath);
+ }
+
+ /**
+ * Writes the accumulated audit records to the audit file.
+ * This method is called after each recordOperation to persist the audit log.
+ * Failures to write audit records are logged but do not throw exceptions
+ * to avoid breaking lock operations.
+ */
+ private void writeAuditFile() {
+ try {
+ // Write the entire buffer content to the audit file
+ // This overwrites the file with all accumulated records
+ String content = auditBuffer.toString();
+ boolean success = storageLockClient.writeObject(auditFilePath, content);
+ if (success) {
+ LOG.debug("Successfully wrote audit records to: {}", auditFilePath);
+ } else {
+ LOG.warn("Failed to write audit records to: {}", auditFilePath);
+ }
+ } catch (Exception e) {
+ LOG.warn("Failed to write audit records to: {}", auditFilePath, e);
+ // Don't throw exception - audit failures should not break lock operations
+ }
+ }
+
+ /**
+ * Closes the audit service. Since all audit records are written after each operation,
+ * no additional cleanup is required during close.
+ *
+ * @throws Exception if there's an error during cleanup (not expected in current implementation)
+ */
+ @Override
+ public synchronized void close() throws Exception {
+ // All audit records are already written after each recordOperation()
+ // No additional writes needed during close
+ LOG.debug("Closed StorageLockProviderAuditService for transaction: {}, owner: {}",
+ transactionStartTime, ownerId);
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/metrics/HoodieLockMetrics.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/metrics/HoodieLockMetrics.java
index 25603d5655c8..fb24e1b69ca5 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/metrics/HoodieLockMetrics.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/metrics/HoodieLockMetrics.java
@@ -18,24 +18,37 @@
package org.apache.hudi.client.transaction.lock.metrics;
+import org.apache.hudi.common.util.HoodieTimer;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.metrics.Metrics;
+
import com.codahale.metrics.Counter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.SlidingWindowReservoir;
import com.codahale.metrics.Timer;
-
-import org.apache.hudi.common.util.HoodieTimer;
-import org.apache.hudi.config.HoodieWriteConfig;
-import org.apache.hudi.metrics.Metrics;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import java.util.concurrent.TimeUnit;
public class HoodieLockMetrics {
+ private static final Logger LOG = LoggerFactory.getLogger(HoodieLockMetrics.class);
public static final String LOCK_ACQUIRE_ATTEMPTS_COUNTER_NAME = "lock.acquire.attempts";
public static final String LOCK_ACQUIRE_SUCCESS_COUNTER_NAME = "lock.acquire.success";
public static final String LOCK_ACQUIRE_FAILURES_COUNTER_NAME = "lock.acquire.failure";
+ public static final String LOCK_RELEASE_SUCCESS_COUNTER_NAME = "lock.release.success";
public static final String LOCK_ACQUIRE_DURATION_TIMER_NAME = "lock.acquire.duration";
public static final String LOCK_REQUEST_LATENCY_TIMER_NAME = "lock.request.latency";
+ public static final String LOCK_ACQUIRED_BY_OTHERS_ERROR_COUNTER_NAME = "lock.acquired.others.error";
+ public static final String LOCK_STATE_UNKNOWN_COUNTER_NAME = "lock.state.unknown";
+ public static final String LOCK_ACQUIRE_PRECONDITION_FAILURE_COUNTER_NAME = "lock.acquire.precondition.failure";
+ public static final String LOCK_PROVIDER_FATAL_ERROR_COUNTER_NAME = "lock.provider.fatal.error";
+ public static final String LOCK_RELEASE_FAILURE_COUNTER_NAME = "lock.release.failure";
+ public static final String LOCK_EXPIRATION_DEADLINE_COUNTER_NAME = "lock.expiration.deadline";
+ public static final String LOCK_DANGLING_COUNTER_NAME = "lock.dangling";
+ public static final String LOCK_INTERRUPTED_COUNTER_NAME = "lock.interrupted";
private final HoodieWriteConfig writeConfig;
private final boolean isMetricsEnabled;
private final int keepLastNtimes = 100;
@@ -44,6 +57,14 @@ public class HoodieLockMetrics {
private transient Counter lockAttempts;
private transient Counter successfulLockAttempts;
private transient Counter failedLockAttempts;
+ private transient Counter lockReleaseSuccess;
+ private transient Counter lockAcquiredByOthersError;
+ private transient Counter lockStateUnknown;
+ private transient Counter lockAcquirePreconditionFailure;
+ private transient Counter lockProviderFatalError;
+ private transient Counter lockReleaseFailure;
+ private transient Counter lockDangling;
+ private transient Counter lockInterrupted;
private transient Timer lockDuration;
private transient Timer lockApiRequestDuration;
private static final Object REGISTRY_LOCK = new Object();
@@ -60,7 +81,14 @@ public class HoodieLockMetrics {
lockAttempts = registry.counter(getMetricsName(LOCK_ACQUIRE_ATTEMPTS_COUNTER_NAME));
successfulLockAttempts = registry.counter(getMetricsName(LOCK_ACQUIRE_SUCCESS_COUNTER_NAME));
failedLockAttempts = registry.counter(getMetricsName(LOCK_ACQUIRE_FAILURES_COUNTER_NAME));
-
+ lockReleaseSuccess = registry.counter(getMetricsName(LOCK_RELEASE_SUCCESS_COUNTER_NAME));
+ lockAcquiredByOthersError = registry.counter(getMetricsName(LOCK_ACQUIRED_BY_OTHERS_ERROR_COUNTER_NAME));
+ lockStateUnknown = registry.counter(getMetricsName(LOCK_STATE_UNKNOWN_COUNTER_NAME));
+ lockAcquirePreconditionFailure = registry.counter(getMetricsName(LOCK_ACQUIRE_PRECONDITION_FAILURE_COUNTER_NAME));
+ lockProviderFatalError = registry.counter(getMetricsName(LOCK_PROVIDER_FATAL_ERROR_COUNTER_NAME));
+ lockReleaseFailure = registry.counter(getMetricsName(LOCK_RELEASE_FAILURE_COUNTER_NAME));
+ lockDangling = registry.counter(getMetricsName(LOCK_DANGLING_COUNTER_NAME));
+ lockInterrupted = registry.counter(getMetricsName(LOCK_INTERRUPTED_COUNTER_NAME));
lockDuration = createTimerForMetrics(registry, LOCK_ACQUIRE_DURATION_TIMER_NAME);
lockApiRequestDuration = createTimerForMetrics(registry, LOCK_REQUEST_LATENCY_TIMER_NAME);
}
@@ -75,9 +103,9 @@ public class HoodieLockMetrics {
if (registry.getMetrics().get(metricName) == null) {
synchronized (REGISTRY_LOCK) {
if (registry.getMetrics().get(metricName) == null) {
- lockDuration = new Timer(new SlidingWindowReservoir(keepLastNtimes));
- registry.register(metricName, lockDuration);
- return lockDuration;
+ Timer timer = new Timer(new SlidingWindowReservoir(keepLastNtimes));
+ registry.register(metricName, timer);
+ return timer;
}
}
}
@@ -90,10 +118,18 @@ public class HoodieLockMetrics {
}
}
+ private static void updateMetric(HoodieTimer timer, Timer metric, String lockName) {
+ Option<Long> durationMs = timer.tryEndTimer();
+ if (durationMs.isPresent()) {
+ metric.update(durationMs.get(), TimeUnit.MILLISECONDS);
+ } else {
+ LOG.info("Unable to get lock {} duration", lockName);
+ }
+ }
+
public void updateLockAcquiredMetric() {
if (isMetricsEnabled) {
- long durationMs = lockApiRequestDurationTimer.endTimer();
- lockApiRequestDuration.update(durationMs, TimeUnit.MILLISECONDS);
+ updateMetric(lockApiRequestDurationTimer, lockApiRequestDuration, "acquired");
lockAttempts.inc();
successfulLockAttempts.inc();
lockDurationTimer.startTimer();
@@ -102,16 +138,68 @@ public class HoodieLockMetrics {
public void updateLockNotAcquiredMetric() {
if (isMetricsEnabled) {
- long durationMs = lockApiRequestDurationTimer.endTimer();
- lockApiRequestDuration.update(durationMs, TimeUnit.MILLISECONDS);
+ updateMetric(lockApiRequestDurationTimer, lockApiRequestDuration, "acquired");
failedLockAttempts.inc();
}
}
public void updateLockHeldTimerMetrics() {
if (isMetricsEnabled && lockDurationTimer != null) {
- long lockDurationInMs = lockDurationTimer.endTimer();
- lockDuration.update(lockDurationInMs, TimeUnit.MILLISECONDS);
+ updateMetric(lockDurationTimer, lockDuration, "held");
+ }
+ }
+
+ public void updateLockReleaseSuccessMetric() {
+ if (isMetricsEnabled) {
+ lockReleaseSuccess.inc();
+ }
+ }
+
+ public void updateLockAcquiredByOthersErrorMetric() {
+ if (isMetricsEnabled) {
+ lockAcquiredByOthersError.inc();
+ }
+ }
+
+ public void updateLockStateUnknownMetric() {
+ if (isMetricsEnabled) {
+ lockStateUnknown.inc();
+ }
+ }
+
+ public void updateLockAcquirePreconditionFailureMetric() {
+ if (isMetricsEnabled) {
+ lockAcquirePreconditionFailure.inc();
+ }
+ }
+
+ public void updateLockProviderFatalErrorMetric() {
+ if (isMetricsEnabled) {
+ lockProviderFatalError.inc();
+ }
+ }
+
+ public void updateLockReleaseFailureMetric() {
+ if (isMetricsEnabled) {
+ lockReleaseFailure.inc();
+ }
+ }
+
+ public void updateLockExpirationDeadlineMetric(int deadlineMs) {
+ if (isMetricsEnabled) {
+ metrics.registerGauge(getMetricsName(LOCK_EXPIRATION_DEADLINE_COUNTER_NAME), deadlineMs >= 0 ? (long)deadlineMs : 0);
+ }
+ }
+
+ public void updateLockDanglingMetric() {
+ if (isMetricsEnabled) {
+ this.lockDangling.inc();
+ }
+ }
+
+ public void updateLockInterruptedMetric() {
+ if (isMetricsEnabled) {
+ this.lockInterrupted.inc();
}
}
}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/HeartbeatManager.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/HeartbeatManager.java
new file mode 100644
index 000000000000..97c9b864d498
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/HeartbeatManager.java
@@ -0,0 +1,49 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+/**
+ * The heartbeat manager interface is meant to manage the lifecycle of heartbeat tasks.
+ *
+ */
+public interface HeartbeatManager extends AutoCloseable {
+
+ /**
+ * Starts the heartbeat for the given thread and does not stop until stopHeartbeat is called or the thread has died.
+ * @param threadToMonitor The thread to pass to/monitor when running the heartbeat task.
+ * @return @return True when there is no previously active heartbeat and the heartbeat is successfully started. False
+ * otherwise.
+ */
+ boolean startHeartbeatForThread(Thread threadToMonitor);
+
+ /**
+ * Stops the heartbeat, if one is active.
+ * This is a blocking call, which drains any in-flight heart beat task execution before return.
+ * @param mayInterruptIfRunning Whether we may interrupt the underlying heartbeat task if it is in-flight.
+ * @return true: no heartbeat task is in-flight or to be executed.
+ * false: failed to stop the heartbeat, there can still be recurring execution of heartbeat tasks.
+ */
+ boolean stopHeartbeat(boolean mayInterruptIfRunning);
+
+ /**
+ * Whether the heartbeat manager has an active heartbeat task currently.
+ * @return A boolean.
+ */
+ boolean hasActiveHeartbeat();
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockGetResult.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockGetResult.java
new file mode 100644
index 000000000000..68809d244c3d
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockGetResult.java
@@ -0,0 +1,38 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+public enum LockGetResult {
+ // Lock file does not exist with code 0
+ NOT_EXISTS(0),
+ // Successfully retrieved the lock file with code 1
+ SUCCESS(1),
+ // Unable to determine lock state due to transient errors with code 2
+ UNKNOWN_ERROR(2);
+
+ private final int code;
+
+ LockGetResult(int code) {
+ this.code = code;
+ }
+
+ public int getCode() {
+ return code;
+ }
+}
\ No newline at end of file
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockProviderHeartbeatManager.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockProviderHeartbeatManager.java
new file mode 100644
index 000000000000..a91193d5dc05
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockProviderHeartbeatManager.java
@@ -0,0 +1,351 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+import org.apache.hudi.common.util.VisibleForTesting;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+/**
+ * LockProviderHeartbeatManager is a helper class which handles the scheduling and stopping of heartbeat
+ * tasks. This is intended for use with the storage based lock provider, which requires
+ * a separate thread to spawn and renew the lock repeatedly.
+ * It should be responsible for the entire lifecycle of the heartbeat task.
+ * Importantly, a new instance should be created for each lock provider.
+ */
+@ThreadSafe
+public class LockProviderHeartbeatManager implements HeartbeatManager {
+ public static long DEFAULT_STOP_HEARTBEAT_TIMEOUT_MS = 15_000L;
+ @GuardedBy("this")
+ private final ScheduledExecutorService scheduler;
+
+ // Constant does not need multi-threading protections.
+ private final String ownerId;
+ private final Logger logger;
+ private final long heartbeatTimeMs;
+
+ /**
+ * Contract for the heartbeat function execution.
+ *
+ * <p>Behavior of the heartbeat manager (consumer):
+ * <ul>
+ * <li>Executes heartBeatFuncToExec every heartbeatTimeMs when:
+ * <ul>
+ * <li>heartBeatFuncToExec returns true</li>
+ * </ul>
+ * </li>
+ * <li>Stops executing heartBeatFuncToExec when:
+ * <ul>
+ * <li>heartBeatFuncToExec returns false</li>
+ * <li>heartBeatFuncToExec throws an exception</li>
+ * <li>heart beat manager calls stopHeartbeat, which will interrupt any inflight execution
+ * and prevent further recurring executions</li>
+ * </ul>
+ * </li>
+ * </ul>
+ *
+ * <p>Requirements for heartBeatFuncToExec implementation:
+ * <ul>
+ * <li>Should perform the logic of renewing lock lease</li>
+ * <li>Should be super light-weight, typically runs within 1 second</li>
+ * <li>Should handle thread interruptions so that the stopHeartbeat function will not wait long for any
+ * inflight execution to complete</li>
+ * <li>Should almost always return true in cases like:
+ * <ul>
+ * <li>Successfully extending the lock lease</li>
+ * <li>Transient failures (network partition, remote service errors) to allow automatic retry</li>
+ * </ul>
+ * </li>
+ * <li>Should return false only in specific cases:
+ * <ul>
+ * <li>When the lock is already expired (no point in extending an expired lock)</li>
+ * <li>When the writer thread does not hold any lock</li>
+ * </ul>
+ * </li>
+ * </ul>
+ *
+ * <p>Warning: Returning false stops all future lock renewal attempts. If the writer thread
+ * is still running, it will execute with a lock that can expire at any time, potentially
+ * leading to corrupted data.
+ */
+ private final Supplier<Boolean> heartbeatFuncToExec;
+ private final long stopHeartbeatTimeoutMs;
+
+ // We ensure within the context of LockProviderHeartbeatManager, heartbeatFuncToExec only execute in a single thread periodically.
+ @GuardedBy("this")
+ private ScheduledFuture<?> scheduledFuture;
+
+ /**
+ * Semaphore for managing heartbeat task execution synchronization.
+ *
+ * <p><strong>IMPORTANT: Thread Safety Warning</strong>
+ * This semaphore is mutually exclusive with {@code synchronized(this)}. Never synchronize
+ * on {@code this} while the thread is holding the heartbeatSemaphore.
+ *
+ * <p>Execution flow:
+ * <ul>
+ * <li>Heartbeat task always attempts to acquire the semaphore before proceeding and
+ * releases it before finishing the current round of execution</li>
+ * <li>The heartbeat manager acquires the semaphore when it needs to:
+ * <ul>
+ * <li>Drain any inflight execution</li>
+ * <li>Prevent further execution of the heartbeat task</li>
+ * </ul>
+ * </li>
+ * </ul>
+ */
+ private final Semaphore heartbeatSemaphore;
+
+ private static final Logger DEFAULT_LOGGER = LoggerFactory.getLogger(LockProviderHeartbeatManager.class);
+
+ /**
+ * Initializes a heartbeat manager.
+ * @param ownerId The identifier for logging of who owns this heartbeat manager.
+ * @param heartbeatTimeMs The time between heartbeat executions.
+ * The first heartbeat will execute after this amount of time elapses.
+ * @param heartbeatFuncToExec The function to execute on each heartbeat. This should handle interrupts.
+ */
+ public LockProviderHeartbeatManager(String ownerId,
+ long heartbeatTimeMs,
+ Supplier<Boolean> heartbeatFuncToExec) {
+ this(
+ ownerId,
+ createThreadScheduler((ownerId != null && ownerId.length() >= 6) ? ownerId.substring(0, 6) : ""),
+ heartbeatTimeMs,
+ DEFAULT_STOP_HEARTBEAT_TIMEOUT_MS,
+ heartbeatFuncToExec,
+ new Semaphore(1),
+ DEFAULT_LOGGER);
+ }
+
+ @VisibleForTesting
+ LockProviderHeartbeatManager(String ownerId,
+ ScheduledExecutorService scheduler,
+ long heartbeatTimeMs,
+ long stopHeartbeatTimeoutMs,
+ Supplier<Boolean> heartbeatFuncToExec,
+ Semaphore heartbeatSemaphore,
+ Logger testLogger) {
+ this.ownerId = ownerId;
+ this.heartbeatTimeMs = heartbeatTimeMs;
+ this.heartbeatFuncToExec = heartbeatFuncToExec;
+ this.logger = testLogger;
+ this.scheduler = scheduler;
+ this.heartbeatSemaphore = heartbeatSemaphore;
+ this.stopHeartbeatTimeoutMs = stopHeartbeatTimeoutMs;
+ }
+
+ /**
+ * Creates a new thread scheduler for heartbeat execution.
+ */
+ private static ScheduledExecutorService createThreadScheduler(String shortUuid) {
+ return Executors.newSingleThreadScheduledExecutor(
+ r -> new Thread(r, "LockProvider-HeartbeatManager-Thread-" + shortUuid));
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ public synchronized boolean startHeartbeatForThread(Thread threadToMonitor) {
+ if (threadToMonitor == null) {
+ throw new IllegalArgumentException("threadToMonitor cannot be null.");
+ }
+
+ if (this.hasActiveHeartbeat()) {
+ logger.warn("Owner {}: Heartbeat is already running.", ownerId);
+ return false;
+ }
+ try {
+ scheduledFuture = scheduler.scheduleAtFixedRate(() -> heartbeatTaskRunner(threadToMonitor), heartbeatTimeMs, heartbeatTimeMs, TimeUnit.MILLISECONDS);
+ logger.debug("Owner {}: Heartbeat started with interval: {} ms", ownerId, heartbeatTimeMs);
+ return true;
+ } catch (Exception e) {
+ logger.error("Owner {}: Unable to schedule heartbeat task. {}", ownerId, e);
+ return false;
+ }
+ }
+
+ /**
+ * Responsible for managing the execution and result of the heartbeat task.
+ * Maintains a semaphore which ensures thread safety for determining the state
+ * of the heartbeat (is the heartbeat executing or not).
+ * @param threadToMonitor The thread to monitor. Required by heartbeat execution.
+ */
+ private void heartbeatTaskRunner(Thread threadToMonitor) {
+ if (!heartbeatSemaphore.tryAcquire()) {
+ logger.error("Owner {}: Heartbeat semaphore should be acquirable at the start of every heartbeat!", ownerId);
+ return;
+ }
+
+ boolean heartbeatExecutionSuccessful;
+ try {
+ heartbeatExecutionSuccessful = executeHeartbeat(threadToMonitor);
+ } finally {
+ heartbeatSemaphore.release();
+ }
+
+ // Call synchronized method after releasing the semaphore
+ if (!heartbeatExecutionSuccessful) {
+ logger.warn("Owner {}: Heartbeat function did not succeed.", ownerId);
+ // Unschedule self from further execution if heartbeat was unsuccessful.
+ heartbeatTaskUnscheduleItself();
+ }
+ }
+
+ /**
+ * Executes the heartbeat task.
+ * @param threadToMonitor The thread to monitor. If we detect that
+ * this thread has stopped we should end the heartbeat.
+ * @return Whether the heartbeat task successfully ran.
+ */
+ private boolean executeHeartbeat(Thread threadToMonitor) {
+ // Check if monitored thread is dead
+ if (!threadToMonitor.isAlive()) {
+ logger.warn("Owner {}: Monitored thread is no longer alive.", ownerId);
+ return false;
+ }
+
+ // Execute heartbeat function
+ try {
+ return heartbeatFuncToExec.get();
+ } catch (Exception e) {
+ logger.error("Owner {}: Heartbeat function threw exception {}", ownerId, e);
+ }
+ return false;
+ }
+
+ /**
+ * This prevents further scheduling of the heartbeat task. Intended to be used by heartbeat task itself.
+ */
+ private synchronized void heartbeatTaskUnscheduleItself() {
+ // Do not interrupt this current task.
+ // This will cancel all future invocations.
+ if (scheduledFuture != null) {
+ boolean cancellationSuccessful = scheduledFuture.cancel(true);
+ logger.info("Owner {}: Requested termination of heartbeat task. Cancellation returned {}.", this.ownerId, cancellationSuccessful);
+ scheduledFuture = null;
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ public boolean stopHeartbeat(boolean mayInterruptIfRunning) {
+ if (cancelRecurringHeartbeatTask(mayInterruptIfRunning)) {
+ return false;
+ }
+
+ // If we requested to stop heartbeat, here we ensure the cancel request results in heartbeat task
+ // exiting synchronously.
+ boolean heartbeatStillInflight = syncWaitInflightHeartbeatTaskToFinish();
+ if (heartbeatStillInflight) {
+ // If waiting for cancellation was interrupted, do not log an error.
+ if (Thread.currentThread().isInterrupted()) {
+ logger.warn("Owner {}: Heartbeat is still in flight due to interruption!", ownerId);
+ } else {
+ logger.error("Owner {}: Heartbeat is still in flight!", ownerId);
+ }
+ return false;
+ }
+
+ // We have stopped the heartbeat, now clean up any leftover states.
+ synchronized (this) {
+ logger.debug("Owner {}: Heartbeat task successfully terminated.", ownerId);
+ scheduledFuture = null;
+ }
+ return true;
+ }
+
+ /**
+ * Cancels the recurring heartbeat task.
+ * @param mayInterruptIfRunning Whether to interrupt the heartbeat task if it is currently running.
+ * @return True if the heartbeat task did not need to be stopped.
+ */
+ private synchronized boolean cancelRecurringHeartbeatTask(boolean mayInterruptIfRunning) {
+ if (!this.hasActiveHeartbeat()) {
+ logger.warn("Owner {}: No active heartbeat task to stop.", ownerId);
+ return true;
+ }
+
+ // Attempt to cancel the scheduled future
+ boolean cancellationSuccessful = scheduledFuture.cancel(mayInterruptIfRunning);
+ logger.debug("Owner {}: Requested termination of heartbeat task. Cancellation returned {}", ownerId, cancellationSuccessful);
+ return false;
+ }
+
+ private boolean syncWaitInflightHeartbeatTaskToFinish() {
+ // Wait for up to stopHeartbeatTimeoutMs for the currently executing heartbeat task to complete.
+ // It is assumed that the heartbeat task, when finishing its execution,
+ // sets heartbeatIsExecuting to false and calls notifyAll() on this object.
+ boolean heartbeatStillInflight = true;
+ try {
+ // Semaphore successfully acquired here excludes the heart execution task. tryAcquire with timeout
+ // means we wait any inflight task execution to finish synchronously.
+ heartbeatStillInflight = !heartbeatSemaphore.tryAcquire(stopHeartbeatTimeoutMs, TimeUnit.MILLISECONDS);
+ if (heartbeatStillInflight) {
+ logger.warn("Owner {}: Timed out while waiting for heartbeat termination.", ownerId);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ logger.warn("Owner {}: Interrupted while waiting for heartbeat termination.", ownerId);
+ }
+ // If we successfully acquired the semaphore before, return it here.
+ heartbeatSemaphore.release(heartbeatStillInflight ? 0 : 1);
+ return heartbeatStillInflight;
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public synchronized boolean hasActiveHeartbeat() {
+ return scheduledFuture != null;
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public synchronized void close() throws Exception {
+ if (hasActiveHeartbeat()) {
+ stopHeartbeat(true);
+ }
+ scheduler.shutdown();
+
+ try {
+ if (!scheduler.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)) {
+ scheduler.shutdownNow();
+ }
+ } catch (InterruptedException e) {
+ scheduler.shutdownNow();
+ Thread.currentThread().interrupt();
+ }
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockUpsertResult.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockUpsertResult.java
new file mode 100644
index 000000000000..eec707d4e56d
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/LockUpsertResult.java
@@ -0,0 +1,38 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+public enum LockUpsertResult {
+ // Lock was successfully created/updated with code 0
+ SUCCESS(0),
+ // Another process has modified the lock file (precondition failure) with code 1
+ ACQUIRED_BY_OTHERS(1),
+ // Unable to determine lock state due to transient errors with code 2
+ UNKNOWN_ERROR(2);
+
+ private final int code;
+
+ LockUpsertResult(int code) {
+ this.code = code;
+ }
+
+ public int getCode() {
+ return code;
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/StorageLockData.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/StorageLockData.java
new file mode 100644
index 000000000000..01c7d141ebda
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/StorageLockData.java
@@ -0,0 +1,71 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * Pojo for conditional writes-based lock provider.
+ */
+public class StorageLockData {
+ private final boolean expired;
+ private final long validUntil;
+ private final String owner;
+
+ /**
+ * Initializes an object describing a conditionally written lock.
+ * @param expired Whether the lock is expired.
+ * @param validUntil The epoch in ms when the lock is expired.
+ * @param owner The uuid owner of the owner of this lock.
+ */
+ @JsonCreator
+ public StorageLockData(
+ @JsonProperty(value = "expired", required = true) boolean expired,
+ @JsonProperty(value = "validUntil", required = true) long validUntil,
+ @JsonProperty(value = "owner", required = true) String owner) {
+ this.expired = expired;
+ this.validUntil = validUntil;
+ this.owner = owner;
+ }
+
+ /**
+ * Gets the expiration.
+ * @return The long representing the expiration in ms.
+ */
+ public long getValidUntil() {
+ return this.validUntil;
+ }
+
+ /**
+ * Whether the lock is expired.
+ * @return True boolean representing whether the lock is expired.
+ */
+ public boolean isExpired() {
+ return this.expired;
+ }
+
+ /**
+ * The owner.
+ * @return A string representing the uuid of the owner of this lock.
+ */
+ public String getOwner() {
+ return this.owner;
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/StorageLockFile.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/StorageLockFile.java
new file mode 100644
index 000000000000..4929d207d9a4
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/lock/models/StorageLockFile.java
@@ -0,0 +1,133 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.common.util.ValidationUtils;
+import org.apache.hudi.exception.HoodieIOException;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+public class StorageLockFile {
+
+ private final StorageLockData data;
+ private final String versionId;
+
+ // Get a custom object mapper. See StorageLockData for required properties.
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper()
+ // This allows us to add new properties down the line.
+ .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
+ // Should not let validUntil or expired be null.
+ .enable(DeserializationFeature.FAIL_ON_NULL_FOR_PRIMITIVES);
+
+ /**
+ * Initializes a StorageLockFile using the data and unique versionId.
+ *
+ * @param data The data in the lock file.
+ * @param versionId The version of this lock file. Used to ensure consistency through conditional writes.
+ */
+ public StorageLockFile(StorageLockData data, String versionId) {
+ ValidationUtils.checkArgument(data != null, "Data must not be null.");
+ ValidationUtils.checkArgument(!StringUtils.isNullOrEmpty(versionId), "VersionId must not be null or empty.");
+ this.data = data;
+ this.versionId = versionId;
+ }
+
+ /**
+ * Factory method to load an input stream into a StorageLockFile.
+ *
+ * @param inputStream The input stream of the JSON content.
+ * @param versionId The unique version identifier for the lock file.
+ * @return A new instance of StorageLockFile.
+ * @throws HoodieIOException If deserialization fails.
+ */
+ public static StorageLockFile createFromStream(InputStream inputStream, String versionId) {
+ try {
+ StorageLockData data = OBJECT_MAPPER.readValue(inputStream, StorageLockData.class);
+ return new StorageLockFile(data, versionId);
+ } catch (IOException e) {
+ throw new HoodieIOException("Failed to deserialize JSON content into StorageLockData", e);
+ }
+ }
+
+ /**
+ * Writes the serialized JSON representation of this object to an output stream.
+ *
+ * @param outputStream The output stream to write the JSON to.
+ * @throws HoodieIOException If serialization fails.
+ */
+ public void writeToStream(OutputStream outputStream) {
+ try {
+ OBJECT_MAPPER.writeValue(outputStream, this.data);
+ } catch (IOException e) {
+ throw new HoodieIOException("Error writing object to JSON output stream", e);
+ }
+ }
+
+ /**
+ * Converts the data to a bytearray. Since we know the payloads will be small this is fine.
+ * @return A byte array.
+ * @throws HoodieIOException If serialization fails.
+ */
+ public static byte[] toByteArray(StorageLockData data) {
+ try {
+ return OBJECT_MAPPER.writeValueAsBytes(data);
+ } catch (JsonProcessingException e) {
+ throw new HoodieIOException("Error writing object to byte array", e);
+ }
+ }
+
+ /**
+ * Gets the version id.
+ * @return A string for the version id.
+ */
+ public String getVersionId() {
+ return this.versionId;
+ }
+
+ /**
+ * Gets the expiration time in ms.
+ * @return A long representing the expiration.
+ */
+ public long getValidUntilMs() {
+ return this.data.getValidUntil();
+ }
+
+ /**
+ * Gets whether the lock is expired.
+ * @return A boolean representing expired.
+ */
+ public boolean isExpired() {
+ return this.data.isExpired();
+ }
+
+ /**
+ * Gets the owner of the lock.
+ * @return A string for the owner of the lock.
+ */
+ public String getOwner() {
+ return this.data.getOwner();
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/StorageBasedLockConfig.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/StorageBasedLockConfig.java
new file mode 100644
index 000000000000..b72450f533c3
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/StorageBasedLockConfig.java
@@ -0,0 +1,101 @@
+/*
+ * 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.hudi.config;
+
+import org.apache.hudi.common.config.ConfigProperty;
+import org.apache.hudi.common.config.HoodieConfig;
+import org.apache.hudi.common.config.LockConfiguration;
+import org.apache.hudi.common.config.TypedProperties;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.hudi.config.HoodieWriteConfig.BASE_PATH;
+
+public class StorageBasedLockConfig extends HoodieConfig {
+ private static final String SINCE_VERSION_1_0_2 = "1.0.2";
+ private static final String STORAGE_BASED_LOCK_PROPERTY_PREFIX = LockConfiguration.LOCK_PREFIX
+ + "storage.";
+
+ public static final ConfigProperty<Long> VALIDITY_TIMEOUT_SECONDS = ConfigProperty
+ .key(STORAGE_BASED_LOCK_PROPERTY_PREFIX + "validity.timeout.secs")
+ .defaultValue(TimeUnit.MINUTES.toSeconds(5))
+ .markAdvanced()
+ .sinceVersion(SINCE_VERSION_1_0_2)
+ .withDocumentation(
+ "For storage-based lock provider, the amount of time in seconds each new lock is valid for. "
+ + "The lock provider will attempt to renew its lock until it successfully extends the lock lease period "
+ + "or the validity timeout is reached.");
+
+ public static final ConfigProperty<Long> RENEW_INTERVAL_SECS = ConfigProperty
+ .key(STORAGE_BASED_LOCK_PROPERTY_PREFIX + "renew.interval.secs")
+ .defaultValue(30L)
+ .markAdvanced()
+ .sinceVersion(SINCE_VERSION_1_0_2)
+ .withAlternatives(STORAGE_BASED_LOCK_PROPERTY_PREFIX + "heartbeat.poll.secs")
+ .withDocumentation(
+ "For storage-based lock provider, the amount of time in seconds to wait before renewing the lock. "
+ + "Defaults to 30 seconds.");
+
+ public long getValiditySeconds() {
+ return getLong(VALIDITY_TIMEOUT_SECONDS);
+ }
+
+ public long getRenewIntervalSecs() {
+ return getLong(RENEW_INTERVAL_SECS);
+ }
+
+ public String getHudiTableBasePath() {
+ return getString(BASE_PATH);
+ }
+
+ public static class Builder {
+ private final StorageBasedLockConfig lockConfig = new StorageBasedLockConfig();
+
+ public StorageBasedLockConfig build() {
+ lockConfig.setDefaults(StorageBasedLockConfig.class.getName());
+ return lockConfig;
+ }
+
+ public StorageBasedLockConfig.Builder fromProperties(TypedProperties props) {
+ lockConfig.getProps().putAll(props);
+ checkRequiredProps();
+ return this;
+ }
+
+ private void checkRequiredProps() {
+ String notExistsMsg = " does not exist!";
+ if (!lockConfig.contains(BASE_PATH)) {
+ throw new IllegalArgumentException(BASE_PATH.key() + notExistsMsg);
+ }
+ if (lockConfig.getLongOrDefault(VALIDITY_TIMEOUT_SECONDS) < lockConfig.getLongOrDefault(RENEW_INTERVAL_SECS)
+ * 10) {
+ throw new IllegalArgumentException(
+ VALIDITY_TIMEOUT_SECONDS.key() + " should be greater than or equal to 10x " + RENEW_INTERVAL_SECS.key());
+ }
+ if (lockConfig.getLongOrDefault(VALIDITY_TIMEOUT_SECONDS) < 10) {
+ throw new IllegalArgumentException(
+ VALIDITY_TIMEOUT_SECONDS.key() + " should be greater than or equal to 10 seconds.");
+ }
+ if (lockConfig.getLongOrDefault(RENEW_INTERVAL_SECS) < 1) {
+ throw new IllegalArgumentException(
+ RENEW_INTERVAL_SECS.key() + " should be greater than or equal to 1 second.");
+ }
+ }
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/MetricsReporterFactory.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/MetricsReporterFactory.java
index 27034735a040..6ffcc13a552b 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/MetricsReporterFactory.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/MetricsReporterFactory.java
@@ -24,7 +24,6 @@ import org.apache.hudi.common.util.StringUtils;
import org.apache.hudi.config.HoodieWriteConfig;
import org.apache.hudi.config.metrics.HoodieMetricsConfig;
import org.apache.hudi.exception.HoodieException;
-import org.apache.hudi.metrics.cloudwatch.CloudWatchMetricsReporter;
import org.apache.hudi.metrics.custom.CustomizableMetricsReporter;
import org.apache.hudi.metrics.datadog.DatadogMetricsReporter;
import org.apache.hudi.metrics.prometheus.PrometheusReporter;
@@ -87,7 +86,8 @@ public class MetricsReporterFactory {
reporter = new ConsoleMetricsReporter(registry);
break;
case CLOUDWATCH:
- reporter = new CloudWatchMetricsReporter(config, registry);
+ reporter = (MetricsReporter) ReflectionUtils.loadClass("org.apache.hudi.aws.metrics.cloudwatch.CloudWatchMetricsReporter",
+ new Class[]{HoodieWriteConfig.class, MetricRegistry.class}, config, registry);
break;
default:
LOG.error("Reporter type[" + type + "] is not supported.");
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/TestHoodieLockMetrics.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/TestHoodieLockMetrics.java
new file mode 100644
index 000000000000..daaf9a97a447
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/TestHoodieLockMetrics.java
@@ -0,0 +1,348 @@
+/*
+ * 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.hudi.client.transaction;
+
+import com.codahale.metrics.Counter;
+import com.codahale.metrics.MetricRegistry;
+import org.apache.hudi.client.transaction.lock.metrics.HoodieLockMetrics;
+import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.config.metrics.HoodieMetricsConfig;
+import org.apache.hudi.metrics.MetricsReporterType;
+import org.apache.hudi.metrics.Metrics;
+
+import com.codahale.metrics.Gauge;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+public class TestHoodieLockMetrics {
+
+ @Test
+ public void testMetricsHappyPath() {
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(HoodieWriteConfig.newBuilder()
+ .forTable("idk").withPath("/dsfasdf/asdf")
+ .withMetricsConfig(metricsConfig)
+ .build());
+
+ //lock acquired
+ assertDoesNotThrow(lockMetrics::startLockApiTimerContext);
+ assertDoesNotThrow(lockMetrics::updateLockAcquiredMetric);
+ assertDoesNotThrow(lockMetrics::updateLockHeldTimerMetrics);
+
+ //lock not acquired
+ assertDoesNotThrow(lockMetrics::startLockApiTimerContext);
+ assertDoesNotThrow(lockMetrics::updateLockNotAcquiredMetric);
+ }
+
+ @Test
+ public void testMetricsMisses() {
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(HoodieWriteConfig.newBuilder()
+ .forTable("idk").withPath("/dsfasdf/asdf")
+ .withMetricsConfig(metricsConfig)
+ .build());
+
+ assertDoesNotThrow(lockMetrics::updateLockHeldTimerMetrics);
+ assertDoesNotThrow(lockMetrics::updateLockNotAcquiredMetric);
+ assertDoesNotThrow(lockMetrics::updateLockAcquiredMetric);
+ }
+
+ @Test
+ public void testLockReleaseSuccessMetric() {
+ // Test that lock release success metric is properly tracked
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieWriteConfig writeConfig = HoodieWriteConfig.newBuilder()
+ .forTable("testTable").withPath("/test/path")
+ .withMetricsConfig(metricsConfig)
+ .build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(writeConfig);
+
+ // Get the metrics registry to verify counter values
+ Metrics metrics = Metrics.getInstance(writeConfig);
+ MetricRegistry registry = metrics.getRegistry();
+
+ // Verify the lock release success counter exists
+ String metricName = writeConfig.getMetricReporterMetricsNamePrefix() + "." + HoodieLockMetrics.LOCK_RELEASE_SUCCESS_COUNTER_NAME;
+ Counter lockReleaseSuccessCounter = registry.getCounters().get(metricName);
+ assertNotNull(lockReleaseSuccessCounter, "Lock release success counter should exist");
+
+ long initialCount = lockReleaseSuccessCounter.getCount();
+
+ // Simulate successful lock release
+ lockMetrics.updateLockReleaseSuccessMetric();
+ assertEquals(initialCount + 1, lockReleaseSuccessCounter.getCount(), "Lock release success counter should increment by 1");
+
+ // Simulate multiple successful lock releases
+ lockMetrics.updateLockReleaseSuccessMetric();
+ lockMetrics.updateLockReleaseSuccessMetric();
+ assertEquals(initialCount + 3, lockReleaseSuccessCounter.getCount(), "Lock release success counter should increment by 3 total");
+ }
+
+ @Test
+ public void testLockLifecycleWithReleaseSuccess() {
+ // Test complete lock lifecycle including acquisition and successful release
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieWriteConfig writeConfig = HoodieWriteConfig.newBuilder()
+ .forTable("testTable").withPath("/test/path")
+ .withMetricsConfig(metricsConfig)
+ .build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(writeConfig);
+
+ // Get the metrics registry to verify counter values
+ Metrics metrics = Metrics.getInstance(writeConfig);
+ MetricRegistry registry = metrics.getRegistry();
+
+ String acquireMetricName = writeConfig.getMetricReporterMetricsNamePrefix() + "." + HoodieLockMetrics.LOCK_ACQUIRE_SUCCESS_COUNTER_NAME;
+ String releaseMetricName = writeConfig.getMetricReporterMetricsNamePrefix() + "." + HoodieLockMetrics.LOCK_RELEASE_SUCCESS_COUNTER_NAME;
+ Counter lockAcquiredCounter = registry.getCounters().get(acquireMetricName);
+ Counter lockReleaseSuccessCounter = registry.getCounters().get(releaseMetricName);
+
+ long initialAcquireCount = lockAcquiredCounter.getCount();
+ long initialReleaseCount = lockReleaseSuccessCounter.getCount();
+
+ // Simulate complete lock lifecycle
+ lockMetrics.startLockApiTimerContext();
+ lockMetrics.updateLockAcquiredMetric();
+ assertEquals(initialAcquireCount + 1, lockAcquiredCounter.getCount(), "Lock acquired counter should increment by 1");
+ assertEquals(initialReleaseCount, lockReleaseSuccessCounter.getCount(), "Lock release success counter should not change yet");
+
+ // Now release the lock successfully
+ lockMetrics.updateLockReleaseSuccessMetric();
+ lockMetrics.updateLockHeldTimerMetrics();
+ assertEquals(initialAcquireCount + 1, lockAcquiredCounter.getCount(), "Lock acquired counter should still be incremented by 1");
+ assertEquals(initialReleaseCount + 1, lockReleaseSuccessCounter.getCount(), "Lock release success counter should increment by 1");
+
+ // Verify metrics balance for multiple cycles
+ for (int i = 0; i < 5; i++) {
+ lockMetrics.startLockApiTimerContext();
+ lockMetrics.updateLockAcquiredMetric();
+ lockMetrics.updateLockReleaseSuccessMetric();
+ lockMetrics.updateLockHeldTimerMetrics();
+ }
+
+ assertEquals(initialAcquireCount + 6, lockAcquiredCounter.getCount(), "Lock acquired counter should increment by 6 total");
+ assertEquals(initialReleaseCount + 6, lockReleaseSuccessCounter.getCount(), "Lock release success counter should increment by 6 total");
+ }
+
+ @Test
+ public void testNewErrorMetrics() {
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(HoodieWriteConfig.newBuilder()
+ .forTable("idk").withPath("/dsfasdf/asdf")
+ .withMetricsConfig(metricsConfig)
+ .build());
+
+ // Test all the new error metrics methods
+ assertDoesNotThrow(lockMetrics::updateLockAcquiredByOthersErrorMetric,
+ "updateLockAcquiredByOthersErrorMetric should not throw");
+ assertDoesNotThrow(lockMetrics::updateLockStateUnknownMetric,
+ "updateLockStateUnknownMetric should not throw");
+ assertDoesNotThrow(lockMetrics::updateLockAcquirePreconditionFailureMetric,
+ "updateLockAcquirePreconditionFailureMetric should not throw");
+ assertDoesNotThrow(lockMetrics::updateLockProviderFatalErrorMetric,
+ "updateLockProviderFatalErrorMetric should not throw");
+ }
+
+ @Test
+ public void testNewErrorMetricsWithDisabledMetrics() {
+ // Test that the new metrics methods work safely when metrics are disabled
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(false).build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(HoodieWriteConfig.newBuilder()
+ .forTable("idk").withPath("/dsfasdf/asdf")
+ .withMetricsConfig(metricsConfig)
+ .build());
+
+ // All methods should be safe to call even when metrics are disabled
+ assertDoesNotThrow(lockMetrics::updateLockAcquiredByOthersErrorMetric,
+ "updateLockAcquiredByOthersErrorMetric should not throw when metrics disabled");
+ assertDoesNotThrow(lockMetrics::updateLockStateUnknownMetric,
+ "updateLockStateUnknownMetric should not throw when metrics disabled");
+ assertDoesNotThrow(lockMetrics::updateLockAcquirePreconditionFailureMetric,
+ "updateLockAcquirePreconditionFailureMetric should not throw when metrics disabled");
+ assertDoesNotThrow(lockMetrics::updateLockProviderFatalErrorMetric,
+ "updateLockProviderFatalErrorMetric should not throw when metrics disabled");
+ }
+
+ @Test
+ public void testCombinedMetricsScenario() {
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(HoodieWriteConfig.newBuilder()
+ .forTable("idk").withPath("/dsfasdf/asdf")
+ .withMetricsConfig(metricsConfig)
+ .build());
+
+ // Test a realistic scenario combining old and new metrics
+ assertDoesNotThrow(() -> {
+ // Start lock acquisition
+ lockMetrics.startLockApiTimerContext();
+
+ // Lock acquisition failed due to unknown state
+ lockMetrics.updateLockStateUnknownMetric();
+ lockMetrics.updateLockNotAcquiredMetric();
+
+ // Retry - start another acquisition attempt
+ lockMetrics.startLockApiTimerContext();
+
+ // This time failed due to precondition failure
+ lockMetrics.updateLockAcquirePreconditionFailureMetric();
+ lockMetrics.updateLockNotAcquiredMetric();
+
+ // Final attempt - lock acquired by others
+ lockMetrics.startLockApiTimerContext();
+ lockMetrics.updateLockAcquiredByOthersErrorMetric();
+ lockMetrics.updateLockNotAcquiredMetric();
+
+ }, "Combined metrics scenario should not throw");
+ }
+
+ @Test
+ public void testLockInterruptedMetric() {
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieWriteConfig writeConfig = HoodieWriteConfig.newBuilder()
+ .forTable("testTable").withPath("/test/path")
+ .withMetricsConfig(metricsConfig)
+ .build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(writeConfig);
+
+ // Get the metrics registry to verify counter values
+ Metrics metrics = Metrics.getInstance(writeConfig);
+ MetricRegistry registry = metrics.getRegistry();
+ String metricName = writeConfig.getMetricReporterMetricsNamePrefix() + "." + HoodieLockMetrics.LOCK_INTERRUPTED_COUNTER_NAME;
+
+ // Test that the interrupted metric can be called
+ assertDoesNotThrow(lockMetrics::updateLockInterruptedMetric,
+ "updateLockInterruptedMetric should not throw");
+
+ // Verify the counter exists and increments
+ Counter interruptedCounter = registry.getCounters().get(metricName);
+ assertNotNull(interruptedCounter, "Lock interrupted counter should exist");
+
+ long initialCount = interruptedCounter.getCount();
+
+ // Call the metric multiple times
+ lockMetrics.updateLockInterruptedMetric();
+ lockMetrics.updateLockInterruptedMetric();
+ lockMetrics.updateLockInterruptedMetric();
+
+ // Verify the counter incremented
+ assertEquals(initialCount + 3, interruptedCounter.getCount(),
+ "Lock interrupted counter should increment by 3");
+ }
+
+ @Test
+ public void testLockExpirationDeadlineGauge() {
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieWriteConfig writeConfig = HoodieWriteConfig.newBuilder()
+ .forTable("testTable").withPath("/test/path")
+ .withMetricsConfig(metricsConfig)
+ .build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(writeConfig);
+
+ // Test that the method doesn't throw
+ assertDoesNotThrow(() -> lockMetrics.updateLockExpirationDeadlineMetric(5000),
+ "updateLockExpirationDeadlineMetric should not throw");
+
+ // Test multiple updates
+ assertDoesNotThrow(() -> {
+ lockMetrics.updateLockExpirationDeadlineMetric(5000);
+ lockMetrics.updateLockExpirationDeadlineMetric(15000);
+ lockMetrics.updateLockExpirationDeadlineMetric(500);
+ lockMetrics.updateLockExpirationDeadlineMetric(-1000);
+ lockMetrics.updateLockExpirationDeadlineMetric(1);
+ }, "Multiple updateLockExpirationDeadlineMetric calls should not throw");
+
+ // Get the metrics registry to verify gauge exists
+ Metrics metrics = Metrics.getInstance(writeConfig);
+ MetricRegistry registry = metrics.getRegistry();
+ String metricName = writeConfig.getMetricReporterMetricsNamePrefix() + "." + HoodieLockMetrics.LOCK_EXPIRATION_DEADLINE_COUNTER_NAME;
+
+ // Verify the gauge exists
+ Gauge<?> deadlineGaugeRaw = registry.getGauges().get(metricName);
+ assertNotNull(deadlineGaugeRaw, "Lock expiration deadline gauge should exist");
+
+ // Test final value (should be 0 from last update)
+ assertEquals(1, ((Number) deadlineGaugeRaw.getValue()).intValue(), "Final gauge value should be 1");
+ }
+
+ @Test
+ public void testLockDanglingMetric() {
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(true).build();
+ HoodieWriteConfig writeConfig = HoodieWriteConfig.newBuilder()
+ .forTable("testTable").withPath("/test/path")
+ .withMetricsConfig(metricsConfig)
+ .build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(writeConfig);
+
+ // Get the metrics registry to verify counter values
+ Metrics metrics = Metrics.getInstance(writeConfig);
+ MetricRegistry registry = metrics.getRegistry();
+ String metricName = writeConfig.getMetricReporterMetricsNamePrefix() + "." + HoodieLockMetrics.LOCK_DANGLING_COUNTER_NAME;
+
+ // Test that the dangling metric can be called
+ assertDoesNotThrow(lockMetrics::updateLockDanglingMetric,
+ "updateLockDanglingMetric should not throw");
+
+ // Verify the counter exists and increments
+ Counter danglingCounter = registry.getCounters().get(metricName);
+ assertNotNull(danglingCounter, "Lock dangling counter should exist");
+
+ long initialCount = danglingCounter.getCount();
+
+ // Call the metric multiple times
+ lockMetrics.updateLockDanglingMetric();
+ lockMetrics.updateLockDanglingMetric();
+
+ // Verify the counter incremented
+ assertEquals(initialCount + 2, danglingCounter.getCount(),
+ "Lock dangling counter should increment by 2");
+ }
+
+ @Test
+ public void testNewMetricsWithDisabledLocking() {
+ // Test that the new metrics methods work safely when locking metrics are disabled
+ HoodieMetricsConfig metricsConfig = HoodieMetricsConfig.newBuilder()
+ .withReporterType(MetricsReporterType.INMEMORY.name()).withLockingMetrics(false).build();
+ HoodieLockMetrics lockMetrics = new HoodieLockMetrics(HoodieWriteConfig.newBuilder()
+ .forTable("testTable").withPath("/test/path")
+ .withMetricsConfig(metricsConfig)
+ .build());
+
+ // All new methods should be safe to call even when locking metrics are disabled
+ assertDoesNotThrow(lockMetrics::updateLockInterruptedMetric,
+ "updateLockInterruptedMetric should not throw when locking metrics disabled");
+ assertDoesNotThrow(() -> lockMetrics.updateLockExpirationDeadlineMetric(5000),
+ "updateLockExpirationDeadlineMetric should not throw when locking metrics disabled");
+ assertDoesNotThrow(lockMetrics::updateLockDanglingMetric,
+ "updateLockDanglingMetric should not throw when locking metrics disabled");
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/TestLockManager.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/TestLockManager.java
index 1b4c08c53299..e22bd89930a1 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/TestLockManager.java
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/TestLockManager.java
@@ -19,6 +19,7 @@
package org.apache.hudi.client.transaction;
import org.apache.hudi.client.transaction.lock.LockManager;
+import org.apache.hudi.client.transaction.lock.StorageBasedLockProvider;
import org.apache.hudi.client.transaction.lock.ZookeeperBasedLockProvider;
import org.apache.hudi.common.model.HoodieFailedWritesCleaningPolicy;
import org.apache.hudi.common.model.WriteConcurrencyMode;
@@ -26,11 +27,13 @@ import org.apache.hudi.common.testutils.HoodieCommonTestHarness;
import org.apache.hudi.config.HoodieCleanConfig;
import org.apache.hudi.config.HoodieLockConfig;
import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.exception.HoodieException;
import org.apache.curator.test.TestingServer;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.Mockito;
@@ -40,6 +43,8 @@ import org.slf4j.LoggerFactory;
import java.io.IOException;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
public class TestLockManager extends HoodieCommonTestHarness {
@@ -119,4 +124,46 @@ public class TestLockManager extends HoodieCommonTestHarness {
.build())
.build();
}
+
+ @Test
+ void testLockManagerTriesMetricsConstructorFirst() {
+ HoodieWriteConfig writeConfig = getStorageBasedLockWriteConfig();
+ LockManager lockManager = new LockManager(writeConfig, this.metaClient.getFs());
+
+ // The getLockProvider() call should fail due to invalid scheme in basePath,
+ // but this validates that we successfully found and tried the metrics constructor
+ // (rather than falling back to the standard constructor due to reflection issues)
+ HoodieException exception = assertThrows(HoodieException.class, lockManager::getLockProvider,
+ "StorageBasedLockProvider should fail with HoodieLockException due to invalid scheme, but this confirms metrics constructor was found");
+
+ // Assert that the exception message references the scheme failure, confirming we got to the
+ // StorageBasedLockProvider instantiation (not a reflection failure)
+ assertTrue(exception.getCause().getCause().getCause().getMessage().startsWith("Unsupported scheme"),
+ "Exception should mention scheme failure, confirming metrics constructor was successfully found and used");
+
+ lockManager.close();
+ }
+
+ @Test
+ void testLockManagerFallbackToStandardConstructor() {
+ // Use ZookeeperBasedLockProvider which doesn't have the metrics constructor
+ HoodieWriteConfig writeConfig = getSingleWriterWriteConfig();
+ LockManager lockManager = new LockManager(writeConfig, this.metaClient.getFs());
+
+ // Get the lock provider to trigger instantiation and fallback
+ assertDoesNotThrow(() -> {
+ lockManager.getLockProvider();
+ });
+
+ lockManager.close();
+ }
+
+ private HoodieWriteConfig getStorageBasedLockWriteConfig() {
+ return HoodieWriteConfig.newBuilder()
+ .withPath(basePath)
+ .withLockConfig(HoodieLockConfig.newBuilder()
+ .withLockProvider(StorageBasedLockProvider.class)
+ .build())
+ .build();
+ }
}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/StorageBasedLockProviderTestBase.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/StorageBasedLockProviderTestBase.java
new file mode 100644
index 000000000000..ec8e96db964f
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/StorageBasedLockProviderTestBase.java
@@ -0,0 +1,290 @@
+/*
+ * 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.hudi.client.transaction.lock;
+
+import org.apache.hudi.common.config.TypedProperties;
+import org.apache.hudi.config.StorageBasedLockConfig;
+import org.apache.hudi.exception.HoodieLockException;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public abstract class StorageBasedLockProviderTestBase {
+
+ protected StorageBasedLockProvider lockProvider;
+ protected static TypedProperties providerProperties;
+ // A method that subclasses must implement to instantiate the correct provider.
+ protected abstract StorageBasedLockProvider createLockProvider();
+
+ @BeforeEach
+ void setUp() {
+ // Create a fresh instance before each test
+ // This will run first
+ providerProperties = new TypedProperties();
+ }
+
+ @AfterEach
+ void tearDown() {
+ // Optionally unlock or close
+ if (lockProvider != null) {
+ lockProvider.unlock();
+ lockProvider.close();
+ }
+ }
+
+ // ----------------------------------------------------
+ // Basic Tests
+ // ----------------------------------------------------
+
+ @Test
+ void testTryLockSuccess() {
+ boolean lockAcquired = lockProvider.tryLock(5, TimeUnit.SECONDS);
+ assertTrue(lockAcquired, "Lock should be successfully acquired within the time limit.");
+ }
+
+ @Test
+ void testTryLockReleasesLock() {
+ assertTrue(lockProvider.tryLock(3, TimeUnit.SECONDS), "Lock should be successfully acquired.");
+ lockProvider.unlock();
+ assertTrue(lockProvider.tryLock(1, TimeUnit.SECONDS), "Lock should be reacquired after being released.");
+ }
+
+ @Test
+ void testTryLockEdgeCaseZeroTimeout() {
+ boolean lockAcquired = lockProvider.tryLock(0, TimeUnit.SECONDS);
+ assertFalse(lockAcquired, "Lock should not be acquired with a zero timeout.");
+ }
+
+ @Test
+ void testLockThreadKilledShouldNotCauseOrphanedHeartbeat() throws InterruptedException {
+ providerProperties.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), 10);
+ providerProperties.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), 1);
+
+ AtomicReference<StorageBasedLockProvider> lp = new AtomicReference<>();
+ // Create a thread with a new lock provider to acquire the lock
+ Thread lockingThread = new Thread(() -> {
+ lp.set(createLockProvider());
+ lp.get().tryLock();
+ });
+ lockingThread.start();
+
+ // Wait for the thread to acquire the lock
+ // Ensure the locking thread is dead.
+ lockingThread.join(2000);
+ assertFalse(lockingThread.isAlive());
+
+ // After the validity expires, should be able to reacquire
+ boolean lockAcquired = lockProvider.tryLock(15, TimeUnit.SECONDS);
+ assertTrue(lockAcquired, "Lock should be reacquired after expiration");
+ assertNotNull(lockProvider.getLock(), "Lock should be reacquired and getLock() should return non-null");
+ lockProvider.unlock();
+ lockProvider.close();
+ lp.get().close();
+ }
+
+ @Test
+ void testTwoLockProvidersCloseAndUnlock() throws InterruptedException {
+ // We had an issue with tryLock(30, TimeUnit.SECONDS) where it was not synchronized
+ // and this created a race condition where closing it would result in improper
+ // execution with the heartbeat manager.
+ StorageBasedLockProvider provider1 = createLockProvider();
+ assertTrue(provider1.tryLock(), "Provider1 should acquire the lock immediately");
+
+ StorageBasedLockProvider provider2 = createLockProvider();
+
+ CountDownLatch finishLatch = new CountDownLatch(1);
+
+ Thread provider2Thread = new Thread(() -> {
+ // We cannot Thread.sleep in a synchronized method, therefore
+ // this will throw while waiting.
+ assertThrows(HoodieLockException.class, () -> provider2.tryLock(10, TimeUnit.SECONDS));
+ finishLatch.countDown();
+ });
+ provider2Thread.start();
+ // We just need the tryLock code path to enter before close
+ provider2.close();
+ provider1.unlock();
+
+ assertTrue(finishLatch.await(5000, TimeUnit.MILLISECONDS));
+ }
+
+ @Test
+ void testLockAfterClosing() {
+ StorageBasedLockProvider provider1 = createLockProvider();
+ assertTrue(provider1.tryLock(), "Provider1 should acquire the lock immediately");
+ provider1.unlock();
+ provider1.close();
+ assertThrows(HoodieLockException.class, provider1::tryLock);
+ }
+
+ @Test
+ void testCloseBeforeUnlocking() {
+ StorageBasedLockProvider provider1 = createLockProvider();
+ assertTrue(provider1.tryLock(), "Provider1 should acquire the lock immediately");
+ provider1.close();
+ provider1.unlock();
+ assertThrows(HoodieLockException.class, provider1::tryLock);
+ }
+
+ @Test
+ void testUnlockWhenNoLockPresent() {
+ lockProvider.unlock();
+ assertNull(lockProvider.getLock());
+ }
+
+ @Test
+ void testTryLockHappyPath() {
+ assertTrue(lockProvider.tryLock(), "tryLock should succeed if lock not held");
+ assertNotNull(lockProvider.getLock(), "Lock should be held after tryLock");
+ }
+
+ @Test
+ void testTryLockReentrancy() {
+ assertTrue(lockProvider.tryLock(), "tryLock should succeed if lock not held");
+ assertTrue(lockProvider.tryLock(), "tryLock should succeed again");
+ }
+
+ @Test
+ void testUnlockAndLock() {
+ assertTrue(lockProvider.tryLock(), "tryLock should succeed if lock not held");
+ assertNotNull(lockProvider.getLock());
+ lockProvider.unlock();
+ assertNull(lockProvider.getLock(), "Lock should be null after unlock");
+ assertTrue(lockProvider.tryLock(), "tryLock should succeed if lock not held");
+ assertNotNull(lockProvider.getLock());
+ }
+
+ @Test
+ void testIdempotentUnlock() {
+ assertTrue(lockProvider.tryLock(), "tryLock should succeed if lock not held");
+ lockProvider.unlock();
+
+ // Calling this again should no op, but will warn.
+ lockProvider.unlock();
+ }
+
+ @Test
+ void testLockReacquisitionInLoop() {
+ for (int i = 0; i < 100; i++) {
+ assertTrue(lockProvider.tryLock(), "tryLock should succeed if lock not held");
+ assertNotNull(lockProvider.getLock(), "Lock should be held after acquisition in iteration " + i);
+ lockProvider.unlock();
+ assertNull(lockProvider.getLock(), "Lock should be null after unlock in iteration " + i);
+ }
+ }
+
+ @Test
+ void testMultipleLockProvidersContention() throws InterruptedException {
+ final int NUM_THREADS = 10;
+ List<StorageBasedLockProvider> providers = new ArrayList<>();
+ List<Thread> threads = new ArrayList<>();
+ AtomicBoolean stopFlag = new AtomicBoolean(false);
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch threadsCompleteLatch = new CountDownLatch(NUM_THREADS);
+
+ // Shared ArrayList that will detect concurrent modifications
+ ArrayList<Integer> sharedList = new ArrayList<>();
+ // Pre-populate the list to allow modifications
+ for (int i = 0; i < 1000; i++) {
+ sharedList.add(i);
+ }
+
+ for (int i = 0; i < NUM_THREADS; i++) {
+ providers.add(createLockProvider());
+ }
+
+ for (int i = 0; i < NUM_THREADS; i++) {
+ final StorageBasedLockProvider provider = providers.get(i);
+ final int threadId = i;
+ Thread t = new Thread(() -> {
+ try {
+ startLatch.await();
+ while (!stopFlag.get()) {
+ if (provider.tryLock()) {
+ // Leveraging java arraylist built in best effort concurrent modification detection
+ // mechanism, if lock failed to process the list, there is a chance of throwing concurrent modification
+ // exception.
+ try {
+ // Multiple operations on ArrayList that must be atomic
+ sharedList.removeIf(val -> val % (threadId + 1) == 0);
+
+ // Add numbers back based on threadId
+ for (int j = 0; j < 1000; j++) {
+ if (j % (threadId + 1) == 0) {
+ sharedList.add(j);
+ }
+ }
+
+ // Sort the list - this will fail if another thread modifies the list
+ sharedList.sort(Integer::compareTo);
+
+ } finally {
+ provider.unlock();
+ }
+ }
+ Thread.yield();
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ threadsCompleteLatch.countDown();
+ }, "LockProviderThread-" + i);
+ threads.add(t);
+ }
+
+ threads.forEach(Thread::start);
+ startLatch.countDown();
+
+ // Run for 3 seconds
+ Thread.sleep(3000);
+ stopFlag.set(true);
+
+ for (Thread t : threads) {
+ t.join();
+ }
+
+ providers.forEach(StorageBasedLockProvider::close);
+ assertTrue(threadsCompleteLatch.await(6, TimeUnit.SECONDS));
+
+ // Validate the invariant of the list.
+ assertEquals(1000, sharedList.size(), "List should contain 1000 elements");
+ for (int i = 0; i < 1000; i++) {
+ assertEquals(i, sharedList.get(i), "List should contain all numbers 0 to 999 in order");
+ }
+ }
+}
\ No newline at end of file
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/TestStorageBasedLockProvider.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/TestStorageBasedLockProvider.java
new file mode 100644
index 000000000000..375741836cd6
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/TestStorageBasedLockProvider.java
@@ -0,0 +1,963 @@
+/*
+ * 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.hudi.client.transaction.lock;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hudi.client.transaction.lock.metrics.HoodieLockMetrics;
+import org.apache.hudi.client.transaction.lock.models.LockUpsertResult;
+import org.apache.hudi.client.transaction.lock.models.StorageLockData;
+import org.apache.hudi.client.transaction.lock.models.StorageLockFile;
+import org.apache.hudi.client.transaction.lock.models.HeartbeatManager;
+import org.apache.hudi.client.transaction.lock.models.LockGetResult;
+import org.apache.hudi.common.config.LockConfiguration;
+import org.apache.hudi.common.config.TypedProperties;
+import org.apache.hudi.common.testutils.HoodieTestUtils;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.config.StorageBasedLockConfig;
+import org.apache.hudi.exception.HoodieLockException;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.slf4j.Logger;
+
+import java.util.Properties;
+import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.hudi.config.HoodieWriteConfig.BASE_PATH;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.fail;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.contains;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.refEq;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Unit test class for StorageBasedLockProvider
+ */
+class TestStorageBasedLockProvider {
+ private StorageBasedLockProvider lockProvider;
+ private StorageLockClient mockLockService;
+ private HeartbeatManager mockHeartbeatManager;
+ private Logger mockLogger;
+ private final String ownerId = UUID.randomUUID().toString();
+ private static final int DEFAULT_LOCK_VALIDITY_MS = 10000;
+
+ @BeforeEach
+ void setupLockProvider() {
+ mockLockService = mock(StorageLockClient.class);
+ mockHeartbeatManager = mock(HeartbeatManager.class);
+ mockLogger = mock(Logger.class);
+ when(mockHeartbeatManager.stopHeartbeat(true)).thenReturn(true);
+ // Mock the readObject method to return Option.empty() to prevent NPE in audit service creation
+ when(mockLockService.readObject(anyString(), anyBoolean())).thenReturn(Option.empty());
+ TypedProperties props = new TypedProperties();
+ props.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "10");
+ props.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "1");
+ props.put(BASE_PATH.key(), "gs://bucket/lake/db/tbl-default");
+
+ lockProvider = spy(new StorageBasedLockProvider(
+ ownerId,
+ props,
+ (a, b, c) -> mockHeartbeatManager,
+ (a, b, c) -> mockLockService,
+ mockLogger,
+ null));
+ }
+
+ @AfterEach
+ void cleanupLockProvider() {
+ lockProvider.close();
+ }
+
+ @Test
+ void testValidLockStorageLocation() {
+ TypedProperties props = new TypedProperties();
+ props.put(BASE_PATH.key(), "s3://bucket/lake/db/tbl-default");
+
+ LockConfiguration lockConf = new LockConfiguration(props);
+ Configuration storageConf = HoodieTestUtils.getDefaultHadoopConf();
+
+ HoodieLockException ex = assertThrows(HoodieLockException.class,
+ () -> new StorageBasedLockProvider(lockConf, storageConf));
+ assertTrue(ex.getMessage().contains("Failed to load and initialize StorageLock"));
+ }
+
+ @Test
+ void testHadoopConstructorWithLockConfigurationAndConfiguration() {
+ TypedProperties props = new TypedProperties();
+ props.put(BASE_PATH.key(), "file:///tmp/lake/db/tbl-default");
+ props.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "20");
+ props.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "2");
+
+ LockConfiguration lockConf = new LockConfiguration(props);
+
+ // Test that constructor throws expected exception when trying to load non-existent storage lock client
+ HoodieLockException ex = assertThrows(HoodieLockException.class,
+ () -> new StorageBasedLockProvider(lockConf, new Configuration()));
+ assertTrue(ex.getMessage().contains("Failed to load and initialize StorageLock"));
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"gs://bucket/lake/db/tbl-default", "s3://bucket/lake/db/tbl-default",
+ "s3a://bucket/lake/db/tbl-default"})
+ void testNonExistentWriteServiceWithDefaults(String tableBasePathString) {
+ TypedProperties props = new TypedProperties();
+ props.put(BASE_PATH.key(), tableBasePathString);
+
+ LockConfiguration lockConf = new LockConfiguration(props);
+ Configuration storageConf = HoodieTestUtils.getDefaultHadoopConf();
+
+ HoodieLockException ex = assertThrows(HoodieLockException.class,
+ () -> new StorageBasedLockProvider(lockConf, storageConf));
+ assertTrue(ex.getMessage().contains("Failed to load and initialize StorageLock"));
+ }
+
+ @Test
+ void testTryLockForTimeUnitThrowsOnInterrupt() throws Exception {
+ doReturn(false).when(lockProvider).tryLock();
+ CountDownLatch latch = new CountDownLatch(1);
+ Thread t = new Thread(() -> {
+ try {
+ lockProvider.tryLock(1, TimeUnit.SECONDS);
+ } catch (HoodieLockException e) {
+ latch.countDown();
+ }
+ });
+ t.start();
+ Thread.sleep(50);
+ t.interrupt();
+ assertTrue(latch.await(2, TimeUnit.SECONDS));
+ }
+
+ @Test
+ void testTryLockForTimeUnitAcquiresLockEventually() throws Exception {
+ AtomicInteger count = new AtomicInteger(0);
+ doAnswer(inv -> count.incrementAndGet() > 2).when(lockProvider).tryLock();
+ CountDownLatch latch = new CountDownLatch(1);
+ Thread t = new Thread(() -> {
+ assertTrue(lockProvider.tryLock(4, TimeUnit.SECONDS));
+ latch.countDown();
+ });
+ t.start();
+ assertTrue(latch.await(5, TimeUnit.SECONDS));
+ }
+
+ @Test
+ void testTryLockForTimeUnitFailsToAcquireLockEventually() throws Exception {
+ AtomicInteger count = new AtomicInteger(0);
+ doAnswer(inv -> count.incrementAndGet() > 2).when(lockProvider).tryLock();
+ CountDownLatch latch = new CountDownLatch(1);
+ Thread t = new Thread(() -> {
+ assertFalse(lockProvider.tryLock(1, TimeUnit.SECONDS));
+ latch.countDown();
+ });
+ t.start();
+ assertTrue(latch.await(2, TimeUnit.SECONDS));
+ }
+
+ @Test
+ void testTryLockSuccess() {
+ long t0 = 1_000L;
+ when(lockProvider.getCurrentEpochMs())
+ .thenReturn(t0);
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, t0 + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(refEq(data), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ boolean acquired = lockProvider.tryLock();
+ assertTrue(acquired);
+ assertEquals(realLockFile, lockProvider.getLock());
+ verify(mockLockService, atLeastOnce()).tryUpsertLockFile(any(), any());
+ }
+
+ @Test
+ void testTryLockSuccessButFailureToStartHeartbeat() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(false);
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(realLockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+
+ boolean acquired = lockProvider.tryLock();
+ assertFalse(acquired);
+ }
+
+ @Test
+ void testTryLockFailsFromOwnerMismatch() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockFile returnedLockFile = new StorageLockFile(
+ new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, "different-owner"), "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(returnedLockFile)));
+
+ HoodieLockException ex = assertThrows(HoodieLockException.class, () -> lockProvider.tryLock());
+ assertTrue(ex.getMessage().contains("Owners do not match"));
+ }
+
+ @Test
+ void testTryLockFailsDueToExistingLock() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS,
+ "other-owner");
+ StorageLockFile existingLock = new StorageLockFile(data, "v2");
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.SUCCESS, Option.of(existingLock)));
+
+ boolean acquired = lockProvider.tryLock();
+ assertFalse(acquired);
+ }
+
+ @Test
+ void testTryLockFailsToUpdateFile() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.ACQUIRED_BY_OTHERS, Option.empty()));
+ assertFalse(lockProvider.tryLock());
+ }
+
+ @Test
+ void testTryLockFailsDueToUnknownState() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.UNKNOWN_ERROR, Option.empty()));
+ assertFalse(lockProvider.tryLock());
+ }
+
+ @Test
+ void testTryLockSucceedsWhenExistingLockExpiredByTime() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() - DEFAULT_LOCK_VALIDITY_MS,
+ "other-owner");
+ StorageLockFile existingLock = new StorageLockFile(data, "v2");
+ StorageLockData newData = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS,
+ ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(newData, "v1");
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.SUCCESS, Option.of(existingLock)));
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(existingLock))))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+ boolean acquired = lockProvider.tryLock();
+ assertTrue(acquired);
+ }
+
+ @Test
+ void testTryLockReentrancySucceeds() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ boolean acquired = lockProvider.tryLock();
+ assertTrue(acquired);
+ // Re-entrancy succeeds
+ assertTrue(lockProvider.tryLock());
+ }
+
+ @Test
+ void testTryLockReentrancyAfterLockExpiredByTime() {
+ // In an extremely unlikely scenario, we could have a local reference to a lock
+ // which is present but expired,
+ // and because we were unable to stop the heartbeat properly, we did not
+ // successfully set it to null.
+ // Due to the nature of the heartbeat manager, this is expected to introduce
+ // some delay, but not be permanently blocking.
+ // There are a few variations of this edge case, so we must test them all.
+
+ // Here the lock is still "unexpired" but the time shows expired.
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() - DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile expiredLock = new StorageLockFile(data, "v1");
+ doReturn(expiredLock).when(lockProvider).getLock();
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData validData = new StorageLockData(false, System.currentTimeMillis() - DEFAULT_LOCK_VALIDITY_MS,
+ ownerId);
+ StorageLockFile validLock = new StorageLockFile(validData, "v2");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(validLock)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ assertTrue(lockProvider.tryLock());
+ }
+
+ @Test
+ void testTryLockReentrancyAfterLockSetExpired() {
+ // In an extremely unlikely scenario, we could have a local reference to a lock
+ // which is present but expired,
+ // and because we were unable to stop the heartbeat properly, we did not
+ // successfully set it to null.
+ // Due to the nature of the heartbeat manager, this is expected to introduce
+ // some delay, but not be permanently blocking.
+ // There are a few variations of this edge case, so we must test them all.
+
+ // Here the lock is "expired" but the time shows unexpired.
+ StorageLockData data = new StorageLockData(true, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile expiredLock = new StorageLockFile(data, "v1");
+ doReturn(expiredLock).when(lockProvider).getLock();
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData validData = new StorageLockData(false, System.currentTimeMillis() - DEFAULT_LOCK_VALIDITY_MS,
+ ownerId);
+ StorageLockFile validLock = new StorageLockFile(validData, "v2");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(validLock)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ assertTrue(lockProvider.tryLock());
+ }
+
+ @Test
+ void testTryLockHeartbeatStillActive() {
+ // In an extremely unlikely scenario, we could have a local reference to a lock
+ // which is present but expired,
+ // and because we were unable to stop the heartbeat properly, we did not
+ // successfully set it to null.
+ // Due to the nature of the heartbeat manager, this is expected to introduce
+ // some delay, but not be permanently blocking.
+ // There are a few variations of this edge case, so we must test them all.
+
+ // Here the heartbeat is still active, so we have to error out.
+ StorageLockData data = new StorageLockData(true, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile expiredLock = new StorageLockFile(data, "v1");
+ doReturn(expiredLock).when(lockProvider).getLock();
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(true);
+ assertThrows(HoodieLockException.class, () -> lockProvider.tryLock());
+ }
+
+ @Test
+ void testUnlockSucceedsAndReentrancy() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+ when(mockHeartbeatManager.stopHeartbeat(true)).thenReturn(true);
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(false);
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(realLockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS,
+ Option.of(new StorageLockFile(new StorageLockData(true, data.getValidUntil(), ownerId), "v2"))));
+ assertTrue(lockProvider.tryLock());
+ when(mockHeartbeatManager.hasActiveHeartbeat())
+ .thenReturn(true) // when we try to stop the heartbeat we will check if heartbeat is active return true.
+ .thenReturn(false); // when try to set lock to expire we will assert no active heartbeat as a precondition.
+ lockProvider.unlock();
+ assertNull(lockProvider.getLock());
+ lockProvider.unlock();
+ }
+
+ @Test
+ void testUnlockFailsToStopHeartbeat() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+ assertTrue(lockProvider.tryLock());
+ when(mockHeartbeatManager.stopHeartbeat(true)).thenReturn(false);
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(true);
+ assertThrows(HoodieLockException.class, () -> lockProvider.unlock());
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(false);
+ }
+
+ @Test
+ void testUnlockThrowsExceptionWhenLockAcquiredByOthers() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+ assertTrue(lockProvider.tryLock());
+
+ when(mockHeartbeatManager.stopHeartbeat(true)).thenReturn(true);
+ when(mockHeartbeatManager.hasActiveHeartbeat())
+ .thenReturn(true)
+ .thenReturn(false);
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(realLockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.ACQUIRED_BY_OTHERS, Option.empty()));
+
+ HoodieLockException exception = assertThrows(HoodieLockException.class, () -> lockProvider.unlock());
+ assertTrue(exception.getMessage().contains("FAILED_TO_RELEASE"));
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(false);
+ }
+
+ @Test
+ void testCloseFailsToStopHeartbeat() {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+ assertTrue(lockProvider.tryLock());
+ when(mockHeartbeatManager.stopHeartbeat(true)).thenReturn(false);
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(true);
+ // Should wrap the exception and log error.
+ lockProvider.close();
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(false);
+ }
+
+ @Test
+ void testRenewLockReturnsFalseWhenNoLockHeld() {
+ doReturn(null).when(lockProvider).getLock();
+ assertFalse(lockProvider.renewLock());
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(true);
+ verify(mockLogger).warn("Owner {}: Cannot renew, no lock held by this process", this.ownerId);
+ }
+
+ @Test
+ void testRenewLockWithoutHoldingLock() {
+ doReturn(null).when(lockProvider).getLock();
+ assertFalse(lockProvider.renewLock());
+ when(mockHeartbeatManager.hasActiveHeartbeat()).thenReturn(false);
+ verify(mockLogger).warn("Owner {}: Cannot renew, no lock held by this process", this.ownerId);
+ }
+
+ @Test
+ void testRenewLockWithFullyExpiredLock() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() - DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile nearExpiredLockFile = new StorageLockFile(data, "v1");
+ doReturn(nearExpiredLockFile).when(lockProvider).getLock();
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(nearExpiredLockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.ACQUIRED_BY_OTHERS, null));
+ assertFalse(lockProvider.renewLock());
+ verify(mockLogger).error("Owner {}: Unable to renew lock as it is acquired by others.", this.ownerId);
+ }
+
+ @Test
+ void testRenewLockUnableToUpsertLockFileButNotFatal() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ doReturn(lockFile).when(lockProvider).getLock();
+ // Signal the upsert attempt failed, but may be transient. See interface for
+ // more details.
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(lockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.UNKNOWN_ERROR, Option.empty()));
+ assertTrue(lockProvider.renewLock());
+ }
+
+ @Test
+ void testRenewLockUnableToUpsertLockFileFatal() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ doReturn(lockFile).when(lockProvider).getLock();
+ // Signal the upsert attempt failed, but may be transient. See interface for
+ // more details.
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(lockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.UNKNOWN_ERROR, Option.empty()));
+ // renewLock return true so it will be retried.
+ assertTrue(lockProvider.renewLock());
+
+ verify(mockLogger).warn("Owner {}: Unable to renew lock due to unknown error, could be transient.", this.ownerId);
+ }
+
+ @Test
+ void testRenewLockSucceedsButRenewalWithinExpirationWindow() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ doReturn(lockFile).when(lockProvider).getLock();
+
+ StorageLockData nearExpirationData = new StorageLockData(false, System.currentTimeMillis(), ownerId);
+ StorageLockFile lockFileNearExpiration = new StorageLockFile(nearExpirationData, "v2");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(lockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(lockFileNearExpiration)));
+
+ // We used to fail in this case before, but since we are only modifying a single
+ // lock file, this is ok now.
+ // Therefore, this can be a happy path variation.
+ assertTrue(lockProvider.renewLock());
+ }
+
+ @Test
+ void testRenewLockSucceeds() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ doReturn(lockFile).when(lockProvider).getLock();
+
+ StorageLockData successData = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS,
+ ownerId);
+ StorageLockFile successLockFile = new StorageLockFile(successData, "v2");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(lockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(successLockFile)));
+ assertTrue(lockProvider.renewLock());
+
+ verify(mockLogger).info(
+ eq("Owner {}: Lock renewal successful. The renewal completes {} ms before expiration for lock {}."),
+ eq(this.ownerId), anyLong(), eq("gs://bucket/lake/db/tbl-default/.hoodie/.locks/table_lock.json"));
+ }
+
+ @Test
+ void testRenewLockFails() {
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ doReturn(lockFile).when(lockProvider).getLock();
+
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.of(lockFile))))
+ .thenThrow(new RuntimeException("Failure"));
+ assertFalse(lockProvider.renewLock());
+
+ verify(mockLogger).error(eq("Owner {}: Exception occurred while renewing lock"), eq(ownerId),
+ any(RuntimeException.class));
+ }
+
+ @Test
+ void testCloseCallsDependencies() throws Exception {
+ lockProvider.close();
+ verify(mockLockService, atLeastOnce()).close();
+ verify(mockHeartbeatManager, atLeastOnce()).close();
+ assertNull(lockProvider.getLock());
+ }
+
+ @Test
+ void testCloseWithErrorForLockService() throws Exception {
+ doThrow(new RuntimeException("Some failure")).when(mockLockService).close();
+ lockProvider.close();
+ verify(mockLogger).error(eq("Owner {}: Lock service failed to close."), eq(ownerId), any(RuntimeException.class));
+ assertNull(lockProvider.getLock());
+ }
+
+ @Test
+ void testCloseWithErrorForHeartbeatManager() throws Exception {
+ doThrow(new RuntimeException("Some failure")).when(mockHeartbeatManager).close();
+ lockProvider.close();
+ verify(mockLogger).error(eq("Owner {}: Heartbeat manager failed to close."), eq(ownerId),
+ any(RuntimeException.class));
+ assertNull(lockProvider.getLock());
+ }
+
+ @Test
+ public void testShutdownHookViaReflection() throws Exception {
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile realLockFile = new StorageLockFile(data, "v1");
+ when(mockLockService.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ boolean acquired = lockProvider.tryLock();
+ assertTrue(acquired);
+ assertEquals(realLockFile, lockProvider.getLock());
+ verify(mockLockService, atLeastOnce()).tryUpsertLockFile(any(), any());
+
+ when(mockLockService.tryUpsertLockFile(any(StorageLockData.class), eq(Option.of(realLockFile))))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(realLockFile)));
+
+ // Mock shutdown
+ lockProvider.shutdown(true);
+
+ // Verify that the expected shutdown behaviors occurred.
+ assertNull(lockProvider.getLock());
+ // We do not execute additional actions
+ verify(mockLockService, never()).close();
+ verify(mockHeartbeatManager, never()).close();
+ }
+
+ @Test
+ public void testShutdownHookWhenNoLockPresent() throws Exception {
+ // Now, when calling shutdown(true), the method should immediately return.
+ lockProvider.shutdown(true);
+
+ // Verify that unlock or close methods are NOT invoked, or adjust expectations accordingly.
+ verify(mockLockService, never()).close();
+ verify(mockHeartbeatManager, never()).close();
+ }
+
+ @Test
+ public void testShutdownHookFailsToBeRemoved() throws Exception {
+ doThrow(new IllegalStateException("Shutdown already in progress")).when(lockProvider).tryRemoveShutdownHook();
+ lockProvider.close();
+ verify(mockLockService, atLeastOnce()).close();
+ verify(mockHeartbeatManager, atLeastOnce()).close();
+ assertNull(lockProvider.getLock());
+ }
+
+ @Test
+ void testShutdownHookFiresDuringTryLockWithTimeout() throws Exception {
+ // This test simulates the scenario where the shutdown hook fires while tryLock(long time, TimeUnit unit)
+ // is in progress, and expects that HoodieLockException is thrown when tryLock is called after shutdown.
+
+ // Setup mocks to simulate lock being held by another owner (to keep tryLock looping)
+ StorageLockData otherOwnerData = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, "other-owner");
+ StorageLockFile otherOwnerLock = new StorageLockFile(otherOwnerData, "v1");
+ when(mockLockService.readCurrentLockFile()).thenReturn(Pair.of(LockGetResult.SUCCESS, Option.of(otherOwnerLock)));
+
+ CountDownLatch tryLockStarted = new CountDownLatch(1);
+ CountDownLatch proceedWithShutdown = new CountDownLatch(1);
+ CountDownLatch shutdownCompleted = new CountDownLatch(1);
+ CountDownLatch tryLockCompleted = new CountDownLatch(1);
+ CountDownLatch exceptionThrown = new CountDownLatch(1);
+
+ // Spy on the real tryLock to know when it's been called and coordinate with shutdown
+ AtomicInteger tryLockCallCount = new AtomicInteger(0);
+ doAnswer(inv -> {
+ int count = tryLockCallCount.incrementAndGet();
+ if (count == 1) {
+ // First call - signal that tryLock has started
+ tryLockStarted.countDown();
+ // Wait for shutdown to be triggered
+ assertTrue(proceedWithShutdown.await(2, TimeUnit.SECONDS));
+ } else {
+ // Subsequent calls - wait briefly for shutdown to complete
+ assertTrue(shutdownCompleted.await(100, TimeUnit.MILLISECONDS));
+ }
+ // Call the real method
+ return inv.callRealMethod();
+ }).when(lockProvider).tryLock();
+
+ // Start a thread that will call tryLock with timeout
+ Thread tryLockThread = new Thread(() -> {
+ try {
+ lockProvider.tryLock(2, TimeUnit.SECONDS);
+ // Should not reach here - exception should be thrown after shutdown
+ fail("Should have thrown HoodieLockException after shutdown");
+ } catch (HoodieLockException e) {
+ // Expected - tryLock should throw exception after shutdown
+ exceptionThrown.countDown();
+ } finally {
+ tryLockCompleted.countDown();
+ }
+ });
+
+ tryLockThread.start();
+
+ // Wait for tryLock to start
+ assertTrue(tryLockStarted.await(2, TimeUnit.SECONDS), "tryLock should have started");
+
+ // Now invoke the shutdown hook while tryLock is in progress
+ // Invoke shutdown in a separate thread to simulate shutdown hook
+ Thread shutdownThread = new Thread(() -> {
+ proceedWithShutdown.countDown(); // Signal tryLock to proceed
+ lockProvider.shutdown(true);
+ shutdownCompleted.countDown(); // Signal that shutdown is complete
+ });
+ shutdownThread.start();
+
+ // Wait for both operations to complete
+ assertTrue(tryLockCompleted.await(5, TimeUnit.SECONDS), "tryLock should complete");
+ assertTrue(exceptionThrown.await(1, TimeUnit.SECONDS), "HoodieLockException should have been thrown");
+ shutdownThread.join(2000);
+
+ // Verify the state after shutdown
+ // The lock should be null after shutdown
+ assertNull(lockProvider.getLock(), "Lock should be null after shutdown hook fires");
+
+ // Verify that tryLock was called at least once
+ verify(lockProvider, atLeastOnce()).tryLock();
+ }
+
+ @Test
+ public void testStorageBasedLockProviderWithMetricsConstructor() {
+ // Create test configuration
+ TypedProperties props = new TypedProperties();
+ props.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "10");
+ props.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "1");
+ props.put(BASE_PATH.key(), "gs://bucket/lake/db/tbl-default");
+ LockConfiguration lockConfiguration = new LockConfiguration(props);
+ Configuration storageConf = HoodieTestUtils.getDefaultHadoopConf();
+
+ // Create a mock HoodieLockMetrics object
+ HoodieLockMetrics mockMetrics = mock(HoodieLockMetrics.class);
+
+ // Test that constructor with metrics works by using the internal constructor to avoid scheme issues
+ StorageBasedLockProvider lockProviderWithMetrics = null;
+ try {
+ // First test that the public constructor with metrics compiles and can be called
+ // We expect this to fail due to lock client instantiation, but it validates the constructor exists
+ assertThrows(Exception.class, () -> {
+ new StorageBasedLockProvider(lockConfiguration, storageConf, mockMetrics);
+ }, "Constructor should exist but fail during lock client instantiation");
+
+ // Now create a working instance using the internal constructor for proper validation
+ lockProviderWithMetrics = new StorageBasedLockProvider(
+ UUID.randomUUID().toString(),
+ props,
+ (a, b, c) -> mock(HeartbeatManager.class),
+ (a, b, c) -> new StubStorageLockClient(a, b, new Properties()),
+ mock(Logger.class),
+ mockMetrics);
+
+ // Verify the lock provider was created successfully
+ assertNotNull(lockProviderWithMetrics, "StorageBasedLockProvider should be created successfully");
+
+ // Verify that it can perform basic operations
+ assertNull(lockProviderWithMetrics.getLock(), "Initially should have no lock");
+
+ } catch (Exception e) {
+ fail("StorageBasedLockProvider creation should not throw unexpected exception: " + e.getMessage());
+ } finally {
+ if (lockProviderWithMetrics != null) {
+ lockProviderWithMetrics.close();
+ }
+ }
+ }
+
+ @Test
+ public void testStorageBasedLockProviderStandardConstructor() {
+ // Create test configuration
+ TypedProperties props = new TypedProperties();
+ props.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "10");
+ props.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "1");
+ props.put(BASE_PATH.key(), "gs://bucket/lake/db/tbl-default");
+ LockConfiguration lockConfiguration = new LockConfiguration(props);
+ Configuration storageConf = HoodieTestUtils.getDefaultHadoopConf();
+
+ // Test that standard constructor works by using the internal constructor to avoid scheme issues
+ StorageBasedLockProvider lockProviderStandard = null;
+ try {
+ // First test that the public standard constructor compiles and can be called
+ // We expect this to fail due to lock client instantiation, but it validates the constructor exists
+ assertThrows(Exception.class, () -> {
+ new StorageBasedLockProvider(lockConfiguration, storageConf);
+ }, "Standard constructor should exist but fail during lock client instantiation");
+
+ // Now create a working instance using the internal constructor for proper validation
+ lockProviderStandard = new StorageBasedLockProvider(
+ UUID.randomUUID().toString(),
+ props,
+ (a, b, c) -> mock(HeartbeatManager.class),
+ (a, b, c) -> new StubStorageLockClient(a, b, new Properties()),
+ mock(Logger.class),
+ null); // No metrics for standard constructor test
+
+ // Verify the lock provider was created successfully
+ assertNotNull(lockProviderStandard, "StorageBasedLockProvider should be created successfully");
+
+ // Verify that it can perform basic operations
+ assertNull(lockProviderStandard.getLock(), "Initially should have no lock");
+
+ } catch (Exception e) {
+ fail("StorageBasedLockProvider creation should not throw unexpected exception: " + e.getMessage());
+ } finally {
+ if (lockProviderStandard != null) {
+ lockProviderStandard.close();
+ }
+ }
+ }
+
+ @Test
+ void testAuditServiceIntegrationWhenConfigNotPresent() {
+ // Test that lock provider works correctly when audit config is not present
+ TypedProperties props = new TypedProperties();
+ props.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "10");
+ props.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "1");
+ props.put(BASE_PATH.key(), "gs://bucket/lake/db/tbl-audit-test");
+
+ // Mock client that returns empty for audit config
+ StorageLockClient auditMockClient = mock(StorageLockClient.class);
+ when(auditMockClient.readObject(anyString(), eq(true)))
+ .thenReturn(Option.empty());
+ when(auditMockClient.readCurrentLockFile())
+ .thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+
+ StorageBasedLockProvider auditLockProvider = new StorageBasedLockProvider(
+ ownerId,
+ props,
+ (a, b, c) -> mockHeartbeatManager,
+ (a, b, c) -> auditMockClient,
+ mockLogger,
+ null);
+
+ // Lock provider should work normally even without audit
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ when(auditMockClient.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(lockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ // tryLock should trigger audit service creation (lazily)
+ assertTrue(auditLockProvider.tryLock());
+
+ // Verify audit config was checked during tryLock
+ verify(auditMockClient, times(1)).readObject(
+ contains(".locks/audit_enabled.json"), eq(true));
+
+ // No audit writes should happen since audit is not present
+ verify(auditMockClient, never()).writeObject(
+ contains(".locks/audit"), anyString());
+
+ auditLockProvider.close();
+ }
+
+ @Test
+ void testAuditServiceIntegrationWhenConfigDisabled() {
+ // Test that lock provider works correctly when audit is explicitly disabled
+ TypedProperties props = new TypedProperties();
+ props.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "10");
+ props.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "1");
+ props.put(BASE_PATH.key(), "gs://bucket/lake/db/tbl-audit-disabled");
+
+ // Mock client that returns disabled config
+ StorageLockClient auditMockClient = mock(StorageLockClient.class);
+ String disabledConfig = "{\"STORAGE_LOCK_AUDIT_SERVICE_ENABLED\": false}";
+ when(auditMockClient.readObject(anyString(), eq(true)))
+ .thenReturn(Option.of(disabledConfig));
+ when(auditMockClient.readCurrentLockFile())
+ .thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+
+ StorageBasedLockProvider auditLockProvider = new StorageBasedLockProvider(
+ ownerId,
+ props,
+ (a, b, c) -> mockHeartbeatManager,
+ (a, b, c) -> auditMockClient,
+ mockLogger,
+ null);
+
+ // Set up lock acquisition
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ when(auditMockClient.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(lockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ // tryLock should trigger audit service creation check
+ assertTrue(auditLockProvider.tryLock());
+
+ // Verify audit config was checked during tryLock
+ verify(auditMockClient, times(1)).readObject(
+ contains(".locks/audit_enabled.json"), eq(true));
+
+ // No audit writes should happen since audit is disabled
+ verify(auditMockClient, never()).writeObject(
+ contains(".locks/audit"), anyString());
+
+ auditLockProvider.close();
+ }
+
+ @Test
+ void testAuditServiceIntegrationWhenConfigEnabled() {
+ // Test that lock provider works correctly when audit is enabled
+ TypedProperties props = new TypedProperties();
+ props.put(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "10");
+ props.put(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "1");
+ props.put(BASE_PATH.key(), "gs://bucket/lake/db/tbl-audit-enabled");
+
+ // Mock client that returns enabled config
+ StorageLockClient auditMockClient = mock(StorageLockClient.class);
+ String enabledConfig = "{\"STORAGE_LOCK_AUDIT_SERVICE_ENABLED\": true}";
+ when(auditMockClient.readObject(anyString(), eq(true)))
+ .thenReturn(Option.of(enabledConfig));
+ when(auditMockClient.readCurrentLockFile())
+ .thenReturn(Pair.of(LockGetResult.NOT_EXISTS, Option.empty()));
+ // Mock writeObject method to return true for audit file writes
+ when(auditMockClient.writeObject(anyString(), anyString()))
+ .thenReturn(true);
+
+ StorageBasedLockProvider auditLockProvider = new StorageBasedLockProvider(
+ ownerId,
+ props,
+ (a, b, c) -> mockHeartbeatManager,
+ (a, b, c) -> auditMockClient,
+ mockLogger,
+ null);
+
+ // Set up lock acquisition
+ StorageLockData data = new StorageLockData(false, System.currentTimeMillis() + DEFAULT_LOCK_VALIDITY_MS, ownerId);
+ StorageLockFile lockFile = new StorageLockFile(data, "v1");
+ when(auditMockClient.tryUpsertLockFile(any(), eq(Option.empty())))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(lockFile)));
+ when(mockHeartbeatManager.startHeartbeatForThread(any())).thenReturn(true);
+
+ // tryLock should trigger audit service creation and START audit
+ assertTrue(auditLockProvider.tryLock());
+
+ // Verify audit config was checked during tryLock
+ verify(auditMockClient, times(1)).readObject(
+ contains(".locks/audit_enabled.json"), eq(true));
+
+ // Verify audit START operation was written
+ verify(auditMockClient, times(1)).writeObject(
+ contains(".locks/audit/"), anyString());
+
+ // Unlock should trigger END audit
+ when(auditMockClient.tryUpsertLockFile(any(), any()))
+ .thenReturn(Pair.of(LockUpsertResult.SUCCESS, Option.of(lockFile)));
+ when(mockHeartbeatManager.stopHeartbeat(anyBoolean())).thenReturn(true);
+ auditLockProvider.unlock();
+
+ // Verify audit END operation was written (total 2 writes: START and END)
+ verify(auditMockClient, times(2)).writeObject(
+ contains(".locks/audit/"), anyString());
+
+ auditLockProvider.close();
+ }
+
+ public static class StubStorageLockClient implements StorageLockClient {
+ public StubStorageLockClient(String ownerId, String lockFileUri, Properties props) {
+ assertTrue(lockFileUri.endsWith("table_lock.json"));
+ }
+
+ @Override
+ public Pair<LockUpsertResult, Option<StorageLockFile>> tryUpsertLockFile(
+ StorageLockData newLockData,
+ Option<StorageLockFile> previousLockFile) {
+ return null;
+ }
+
+ @Override
+ public Pair<LockGetResult, Option<StorageLockFile>> readCurrentLockFile() {
+ return null;
+ }
+
+ @Override
+ public Option<String> readObject(String filePath, boolean checkExistsFirst) {
+ // Stub implementation for testing
+ return Option.empty();
+ }
+
+ @Override
+ public boolean writeObject(String filePath, String content) {
+ // Stub implementation for testing
+ return true;
+ }
+
+ @Override
+ public void close() throws Exception {
+ // stub, no-op
+ }
+ }
+
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/audit/TestAuditServiceFactory.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/audit/TestAuditServiceFactory.java
new file mode 100644
index 000000000000..a53c829e8d75
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/audit/TestAuditServiceFactory.java
@@ -0,0 +1,160 @@
+/*
+ * 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.hudi.client.transaction.lock.audit;
+
+import org.apache.hudi.client.transaction.lock.StorageLockClient;
+import org.apache.hudi.common.util.Option;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Test class for AuditServiceFactory
+ */
+public class TestAuditServiceFactory {
+
+ private StorageLockClient mockStorageLockClient;
+ private final String ownerId = "test-owner";
+ private final String basePath = "s3://bucket/path/to/table";
+
+ @BeforeEach
+ void setUp() {
+ mockStorageLockClient = mock(StorageLockClient.class);
+ }
+
+ @Test
+ void testCreateAuditServiceWhenConfigNotFound() {
+ // Config file doesn't exist
+ String expectedPath = basePath + "/.hoodie/.locks/audit_enabled.json";
+ when(mockStorageLockClient.readObject(eq(expectedPath), eq(true)))
+ .thenReturn(Option.empty());
+
+ Option<AuditService> result = AuditServiceFactory.createLockProviderAuditService(
+ ownerId, basePath, mockStorageLockClient,
+ System.currentTimeMillis(),
+ timestamp -> timestamp + 10000,
+ () -> true);
+
+ // Should return empty for now (no concrete implementation yet)
+ // When implementation is added, this should return a present Option
+ assertTrue(result.isEmpty());
+ verify(mockStorageLockClient).readObject(expectedPath, true);
+ }
+
+ @Test
+ void testCreateAuditServiceWhenDisabledInConfig() {
+ // Config file exists but audit is disabled
+ String expectedPath = basePath + "/.hoodie/.locks/audit_enabled.json";
+ String configJson = "{\"STORAGE_LOCK_AUDIT_SERVICE_ENABLED\": false}";
+ when(mockStorageLockClient.readObject(eq(expectedPath), eq(true)))
+ .thenReturn(Option.of(configJson));
+
+ Option<AuditService> result = AuditServiceFactory.createLockProviderAuditService(
+ ownerId, basePath, mockStorageLockClient,
+ System.currentTimeMillis(),
+ timestamp -> timestamp + 10000,
+ () -> true);
+
+ // Should return empty when audit is disabled
+ assertTrue(result.isEmpty());
+ verify(mockStorageLockClient).readObject(expectedPath, true);
+ }
+
+ @Test
+ void testCreateAuditServiceWhenEnabledInConfig() {
+ // Config file exists and audit is enabled
+ String expectedPath = basePath + "/.hoodie/.locks/audit_enabled.json";
+ String configJson = "{\"STORAGE_LOCK_AUDIT_SERVICE_ENABLED\": true}";
+ when(mockStorageLockClient.readObject(eq(expectedPath), eq(true)))
+ .thenReturn(Option.of(configJson));
+ // Mock writeObject method to return true for audit file writes
+ when(mockStorageLockClient.writeObject(anyString(), anyString()))
+ .thenReturn(true);
+ Option<AuditService> result = AuditServiceFactory.createLockProviderAuditService(
+ ownerId, basePath, mockStorageLockClient,
+ System.currentTimeMillis(),
+ timestamp -> timestamp + 10000, // lockExpirationFunction
+ () -> true); // lockHeldSupplier
+
+ assertTrue(result.isPresent());
+ verify(mockStorageLockClient).readObject(expectedPath, true);
+ }
+
+ @Test
+ void testCreateAuditServiceWithMalformedJson() {
+ // Config file exists but contains invalid JSON
+ String expectedPath = basePath + "/.hoodie/.locks/audit_enabled.json";
+ String malformedJson = "{invalid json}";
+ when(mockStorageLockClient.readObject(eq(expectedPath), eq(true)))
+ .thenReturn(Option.of(malformedJson));
+
+ Option<AuditService> result = AuditServiceFactory.createLockProviderAuditService(
+ ownerId, basePath, mockStorageLockClient,
+ System.currentTimeMillis(),
+ timestamp -> timestamp + 10000,
+ () -> true);
+
+ // Should return empty when JSON is malformed
+ assertTrue(result.isEmpty());
+ verify(mockStorageLockClient).readObject(expectedPath, true);
+ }
+
+ @Test
+ void testCreateAuditServiceWithEmptyConfig() {
+ // Config file exists but doesn't contain the expected field
+ String expectedPath = basePath + "/.hoodie/.locks/audit_enabled.json";
+ String configJson = "{\"some_other_field\": true}";
+ when(mockStorageLockClient.readObject(eq(expectedPath), eq(true)))
+ .thenReturn(Option.of(configJson));
+
+ Option<AuditService> result = AuditServiceFactory.createLockProviderAuditService(
+ ownerId, basePath, mockStorageLockClient,
+ System.currentTimeMillis(),
+ timestamp -> timestamp + 10000,
+ () -> true);
+
+ // Should return empty when field is missing (defaults to false)
+ assertTrue(result.isEmpty());
+ verify(mockStorageLockClient).readObject(expectedPath, true);
+ }
+
+ @Test
+ void testCreateAuditServiceUsesCheckExistsFirst() {
+ // Verify that the factory passes true for checkExistsFirst
+ String expectedPath = basePath + "/.hoodie/.locks/audit_enabled.json";
+ when(mockStorageLockClient.readObject(eq(expectedPath), eq(true)))
+ .thenReturn(Option.empty());
+
+ AuditServiceFactory.createLockProviderAuditService(
+ ownerId, basePath, mockStorageLockClient,
+ System.currentTimeMillis(),
+ timestamp -> timestamp + 10000,
+ () -> true);
+
+ // Should pass true for checkExistsFirst since audit config is rarely present
+ verify(mockStorageLockClient).readObject(expectedPath, true);
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/audit/TestStorageLockProviderAuditService.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/audit/TestStorageLockProviderAuditService.java
new file mode 100644
index 000000000000..f93cd191efe9
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/audit/TestStorageLockProviderAuditService.java
@@ -0,0 +1,380 @@
+/*
+ * 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.hudi.client.transaction.lock.audit;
+
+import org.apache.hudi.client.transaction.lock.StorageLockClient;
+
+import org.apache.hadoop.fs.Path;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.mockito.ArgumentCaptor;
+
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Unit tests for StorageLockProviderAuditService.
+ */
+public class TestStorageLockProviderAuditService {
+
+ private static final String BASE_PATH = "s3://bucket/table";
+ private static final String OWNER_ID = "writer-12345678-9abc-def0-1234-567890abcdef";
+ private static final long TRANSACTION_START_TIME = 1234567890000L;
+ private static final long LOCK_EXPIRATION = 1000000L;
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ private StorageLockClient storageLockClient;
+ private Supplier<Boolean> lockHeldSupplier;
+ private StorageLockProviderAuditService auditService;
+
+ @BeforeEach
+ void setUp() {
+ storageLockClient = mock(StorageLockClient.class);
+ lockHeldSupplier = mock(Supplier.class);
+
+ when(lockHeldSupplier.get()).thenReturn(true);
+
+ auditService = new StorageLockProviderAuditService(
+ BASE_PATH,
+ OWNER_ID,
+ TRANSACTION_START_TIME,
+ storageLockClient,
+ timestamp -> LOCK_EXPIRATION,
+ lockHeldSupplier);
+ }
+
+ /**
+ * Helper method to validate audit record structure and content.
+ *
+ * @param auditRecord The audit record to validate
+ * @param expectedState Expected operation state
+ * @param expectedTimestamp Expected timestamp
+ * @param expectedLockHeld Expected lock held status
+ */
+ private void validateAuditRecord(Map<String, Object> auditRecord, AuditOperationState expectedState,
+ long expectedTimestamp, boolean expectedLockHeld) {
+ assertNotNull(auditRecord, "Audit record should not be null");
+ assertEquals(6, auditRecord.size(), "Audit record should contain exactly 6 fields");
+
+ // Validate all required fields are present
+ assertNotNull(auditRecord.get("ownerId"), "ownerId should be present");
+ assertNotNull(auditRecord.get("transactionStartTime"), "transactionStartTime should be present");
+ assertNotNull(auditRecord.get("timestamp"), "timestamp should be present");
+ assertNotNull(auditRecord.get("state"), "state should be present");
+ assertNotNull(auditRecord.get("lockExpiration"), "lockExpiration should be present");
+ assertNotNull(auditRecord.get("lockHeld"), "lockHeld should be present");
+
+ // Validate field values
+ assertEquals(OWNER_ID, auditRecord.get("ownerId"));
+ assertEquals(TRANSACTION_START_TIME, ((Number) auditRecord.get("transactionStartTime")).longValue());
+ assertEquals(expectedTimestamp, ((Number) auditRecord.get("timestamp")).longValue());
+ assertEquals(expectedState.name(), auditRecord.get("state"));
+ assertEquals(LOCK_EXPIRATION, ((Number) auditRecord.get("lockExpiration")).longValue());
+ assertEquals(expectedLockHeld, auditRecord.get("lockHeld"));
+ }
+
+ /**
+ * Helper method to verify expected file path generation.
+ *
+ * @param actualPath The actual file path captured
+ * @param ownerId The owner ID used
+ * @param transactionStartTime The transaction start time used
+ */
+ private void validateFilePath(String actualPath, String ownerId, long transactionStartTime) {
+ String expectedPath = String.format("%s%s.hoodie%s.locks%saudit%s%d_%s.jsonl",
+ BASE_PATH,
+ Path.SEPARATOR,
+ Path.SEPARATOR,
+ Path.SEPARATOR,
+ Path.SEPARATOR,
+ transactionStartTime,
+ ownerId);
+ assertEquals(expectedPath, actualPath);
+ }
+
+ /**
+ * Comprehensive test for complete audit lifecycle: START -> RENEW -> END.
+ * This consolidates individual operation tests into a single comprehensive test
+ * that validates the entire audit flow.
+ */
+ @Test
+ void testCompleteAuditLifecycle() throws Exception {
+ when(storageLockClient.writeObject(anyString(), anyString())).thenReturn(true);
+
+ long startTime = System.currentTimeMillis();
+ long renewTime = startTime + 1000;
+ long endTime = startTime + 2000;
+
+ // Configure lock held status for each operation
+ when(lockHeldSupplier.get()).thenReturn(true, true, false);
+
+ // Execute complete lifecycle
+ auditService.recordOperation(AuditOperationState.START, startTime);
+ auditService.recordOperation(AuditOperationState.RENEW, renewTime);
+ auditService.recordOperation(AuditOperationState.END, endTime);
+
+ // Capture all write calls
+ ArgumentCaptor<String> pathCaptor = ArgumentCaptor.forClass(String.class);
+ ArgumentCaptor<String> contentCaptor = ArgumentCaptor.forClass(String.class);
+ verify(storageLockClient, times(3)).writeObject(pathCaptor.capture(), contentCaptor.capture());
+
+ // Validate file path (should be consistent across all operations)
+ validateFilePath(pathCaptor.getValue(), OWNER_ID, TRANSACTION_START_TIME);
+
+ // Parse final content - should contain all three records
+ String finalContent = contentCaptor.getValue();
+ String[] lines = finalContent.trim().split("\n");
+ assertEquals(3, lines.length, "Should have three JSON lines");
+
+ // Validate each audit record
+ @SuppressWarnings("unchecked")
+ Map<String, Object> startRecord = OBJECT_MAPPER.readValue(lines[0], Map.class);
+ @SuppressWarnings("unchecked")
+ Map<String, Object> renewRecord = OBJECT_MAPPER.readValue(lines[1], Map.class);
+ @SuppressWarnings("unchecked")
+ Map<String, Object> endRecord = OBJECT_MAPPER.readValue(lines[2], Map.class);
+
+ validateAuditRecord(startRecord, AuditOperationState.START, startTime, true);
+ validateAuditRecord(renewRecord, AuditOperationState.RENEW, renewTime, true);
+ validateAuditRecord(endRecord, AuditOperationState.END, endTime, false);
+ }
+
+ /**
+ * Test data provider for different owner ID formats.
+ *
+ * @return Stream of owner ID test cases
+ */
+ static java.util.stream.Stream<Arguments> ownerIdTestCases() {
+ return java.util.stream.Stream.of(
+ Arguments.of("12345678-9abc-def0-1234-567890abcdef", "Full UUID format"),
+ Arguments.of("abc123", "Short owner ID"),
+ Arguments.of("regular-owner-id-without-uuid", "Regular owner ID"),
+ Arguments.of("writer-12345678-9abc-def0-1234-567890abcdef", "Writer with UUID")
+ );
+ }
+
+ /**
+ * Test audit file naming with different owner ID formats.
+ * Consolidated test that verifies file path generation for various owner ID formats.
+ *
+ * @param ownerId The owner ID to test
+ * @param description Description of the test case
+ */
+ @ParameterizedTest(name = "{1}: {0}")
+ @MethodSource("ownerIdTestCases")
+ void testFileNameWithDifferentOwnerIds(String ownerId, String description) throws Exception {
+ long txnStartTime = System.currentTimeMillis();
+ StorageLockProviderAuditService service = new StorageLockProviderAuditService(
+ BASE_PATH,
+ ownerId,
+ txnStartTime,
+ storageLockClient,
+ timestamp -> LOCK_EXPIRATION,
+ lockHeldSupplier);
+
+ when(storageLockClient.writeObject(anyString(), anyString())).thenReturn(true);
+ service.recordOperation(AuditOperationState.START, System.currentTimeMillis());
+
+ ArgumentCaptor<String> pathCaptor = ArgumentCaptor.forClass(String.class);
+ verify(storageLockClient, times(1)).writeObject(pathCaptor.capture(), anyString());
+
+ validateFilePath(pathCaptor.getValue(), ownerId, txnStartTime);
+ }
+
+ /**
+ * Test data provider for write failure scenarios.
+ *
+ * @return Stream of write failure test cases
+ */
+ static Stream<Arguments> writeFailureTestCases() {
+ return Stream.of(
+ Arguments.of(false, "Write returns false"),
+ Arguments.of(new RuntimeException("Storage error"), "Write throws exception")
+ );
+ }
+
+ /**
+ * Test handling of write failures using parameterized test.
+ * Consolidates multiple write failure scenarios into a single parameterized test.
+ *
+ * @param failureCondition The failure condition (Boolean false or Exception)
+ * @param description Description of the test case
+ */
+ @ParameterizedTest(name = "{1}")
+ @MethodSource("writeFailureTestCases")
+ void testWriteFailureHandling(Object failureCondition, String description) throws Exception {
+ long timestamp = System.currentTimeMillis();
+
+ if (failureCondition instanceof Boolean) {
+ when(storageLockClient.writeObject(anyString(), anyString())).thenReturn((Boolean) failureCondition);
+ } else if (failureCondition instanceof Exception) {
+ when(storageLockClient.writeObject(anyString(), anyString())).thenThrow((Exception) failureCondition);
+ }
+
+ // Should not throw exception - audit failures should be handled gracefully
+ auditService.recordOperation(AuditOperationState.START, timestamp);
+
+ verify(storageLockClient, times(1)).writeObject(anyString(), anyString());
+ }
+
+ @Test
+ void testDynamicFunctionValues() throws Exception {
+ long timestamp = System.currentTimeMillis();
+ when(storageLockClient.writeObject(anyString(), anyString())).thenReturn(true);
+
+ // Create a service with dynamic expiration calculation
+ StorageLockProviderAuditService dynamicService = new StorageLockProviderAuditService(
+ BASE_PATH,
+ OWNER_ID,
+ TRANSACTION_START_TIME,
+ storageLockClient,
+ ts -> ts + 1000L, // Adds 1000ms to timestamp
+ lockHeldSupplier);
+
+ when(lockHeldSupplier.get()).thenReturn(true, false);
+
+ dynamicService.recordOperation(AuditOperationState.START, timestamp);
+ dynamicService.recordOperation(AuditOperationState.END, timestamp + 1000);
+
+ ArgumentCaptor<String> contentCaptor = ArgumentCaptor.forClass(String.class);
+ verify(storageLockClient, times(2)).writeObject(anyString(), contentCaptor.capture());
+
+ // Validate dynamic expiration calculation
+ String[] firstLines = contentCaptor.getAllValues().get(0).trim().split("\n");
+ @SuppressWarnings("unchecked")
+ Map<String, Object> firstRecord = OBJECT_MAPPER.readValue(firstLines[0], Map.class);
+
+ String[] secondLines = contentCaptor.getAllValues().get(1).trim().split("\n");
+ @SuppressWarnings("unchecked")
+ Map<String, Object> secondRecord = OBJECT_MAPPER.readValue(secondLines[1], Map.class);
+
+ assertTrue((Boolean) firstRecord.get("lockHeld"));
+ assertEquals(timestamp + 1000L, ((Number) firstRecord.get("lockExpiration")).longValue());
+
+ assertFalse((Boolean) secondRecord.get("lockHeld"));
+ assertEquals(timestamp + 1000L + 1000L, ((Number) secondRecord.get("lockExpiration")).longValue());
+ }
+
+ @Test
+ void testFilePathGeneration() throws Exception {
+ when(storageLockClient.writeObject(anyString(), anyString())).thenReturn(true);
+
+ auditService.recordOperation(AuditOperationState.START, System.currentTimeMillis());
+
+ String expectedPath = String.format("%s/.hoodie/.locks/audit/%d_%s.jsonl",
+ BASE_PATH, TRANSACTION_START_TIME, OWNER_ID);
+
+ verify(storageLockClient).writeObject(eq(expectedPath), anyString());
+ }
+
+ @Test
+ void testCloseMethodWithBufferedData() throws Exception {
+ when(storageLockClient.writeObject(anyString(), anyString())).thenReturn(true);
+
+ // Record an operation to buffer data
+ auditService.recordOperation(AuditOperationState.START, System.currentTimeMillis());
+
+ // Close should not write additional data (data is written immediately after each operation)
+ auditService.close();
+
+ // Verify write was called only once for the START operation
+ verify(storageLockClient, times(1)).writeObject(anyString(), anyString());
+ }
+
+ @Test
+ void testConcurrentOperations() throws Exception {
+ when(storageLockClient.writeObject(anyString(), anyString())).thenReturn(true);
+ when(lockHeldSupplier.get()).thenReturn(true, false);
+
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch completionLatch = new CountDownLatch(2);
+
+ long baseTime = System.currentTimeMillis();
+
+ // Concurrent RENEW and END operations
+ executor.submit(() -> {
+ try {
+ startLatch.await();
+ auditService.recordOperation(AuditOperationState.RENEW, baseTime + 1000);
+ } catch (Exception e) {
+ // Handle gracefully
+ } finally {
+ completionLatch.countDown();
+ }
+ });
+
+ executor.submit(() -> {
+ try {
+ startLatch.await();
+ auditService.recordOperation(AuditOperationState.END, baseTime + 1001);
+ } catch (Exception e) {
+ // Handle gracefully
+ } finally {
+ completionLatch.countDown();
+ }
+ });
+
+ startLatch.countDown();
+ assertTrue(completionLatch.await(5, TimeUnit.SECONDS));
+
+ // Verify both operations were recorded
+ ArgumentCaptor<String> contentCaptor = ArgumentCaptor.forClass(String.class);
+ verify(storageLockClient, times(2)).writeObject(anyString(), contentCaptor.capture());
+
+ String finalContent = contentCaptor.getValue();
+ String[] lines = finalContent.trim().split("\\n");
+ assertEquals(2, lines.length);
+
+ // Verify both RENEW and END states are present
+ @SuppressWarnings("unchecked")
+ Map<String, Object> record1 = OBJECT_MAPPER.readValue(lines[0], Map.class);
+ @SuppressWarnings("unchecked")
+ Map<String, Object> record2 = OBJECT_MAPPER.readValue(lines[1], Map.class);
+
+ String state1 = (String) record1.get("state");
+ String state2 = (String) record2.get("state");
+ assertTrue((state1.equals("RENEW") && state2.equals("END"))
+ || (state1.equals("END") && state2.equals("RENEW")));
+
+ executor.shutdown();
+ assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS));
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/models/StorageLockClientFileTest.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/models/StorageLockClientFileTest.java
new file mode 100644
index 000000000000..9a415850289a
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/models/StorageLockClientFileTest.java
@@ -0,0 +1,223 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+import org.apache.hudi.client.transaction.lock.StorageLockClient;
+import org.apache.hudi.client.transaction.lock.audit.StorageLockProviderAuditService;
+import org.apache.hudi.exception.HoodieIOException;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+
+public class StorageLockClientFileTest {
+
+ private static final String JSON_DATA = "{\"expired\":false,\"validUntil\":1700000000000,\"owner\":\"testOwner\"}";
+ private static final String JSON_DATA_EXTRA_FIELD = "{\"expired\":true,\"validUntil\":1600000000000,\"owner\":\"otherOwner\",\"state\":\"active\"}";
+ private static final String INVALID_JSON = "{\"invalidField\":123}";
+ private static final String VERSION_ID = "testVersionId";
+
+ private InputStream validJsonStream;
+ private InputStream extraFieldValidJsonStream;
+ private InputStream invalidJsonStream;
+
+ @BeforeEach
+ void setup() {
+ validJsonStream = new ByteArrayInputStream(JSON_DATA.getBytes());
+ extraFieldValidJsonStream = new ByteArrayInputStream(JSON_DATA_EXTRA_FIELD.getBytes());
+ invalidJsonStream = new ByteArrayInputStream(INVALID_JSON.getBytes());
+ }
+
+ @Test
+ void testCreateValidInputStream() {
+ StorageLockFile file = StorageLockFile.createFromStream(validJsonStream, VERSION_ID);
+ assertEquals(1700000000000L, file.getValidUntilMs());
+ assertEquals("testOwner", file.getOwner());
+ assertEquals(VERSION_ID, file.getVersionId());
+ assertFalse(file.isExpired());
+ }
+
+ @Test
+ void testCreateValidInputStreamExtraField() {
+ StorageLockFile file = StorageLockFile.createFromStream(extraFieldValidJsonStream, VERSION_ID);
+ assertEquals(1600000000000L, file.getValidUntilMs());
+ assertEquals("otherOwner", file.getOwner());
+ assertEquals(VERSION_ID, file.getVersionId());
+ assertTrue(file.isExpired());
+ }
+
+ @Test
+ void testCreateInvalidInputStreamFromMock() throws IOException {
+ InputStream mockInputStream = mock(InputStream.class);
+
+ doThrow(new IOException("Simulated IOException"))
+ .when(mockInputStream)
+ .read();
+ HoodieIOException exception = assertThrows(HoodieIOException.class, () -> StorageLockFile.createFromStream(mockInputStream, "versionId"));
+ assertTrue(exception.getMessage().contains("Failed to deserialize"));
+ }
+
+ @Test
+ void testCreateInvalidInputStreamFromBadData() {
+ HoodieIOException exception = assertThrows(HoodieIOException.class, () ->
+ StorageLockFile.createFromStream(invalidJsonStream, VERSION_ID)
+ );
+ assertTrue(exception.getMessage().contains("Failed to deserialize"));
+ }
+
+ @Test
+ void testCreateNullData() {
+ IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () ->
+ new StorageLockFile(null, VERSION_ID)
+ );
+ assertTrue(exception.getMessage().contains("Data must not be null"));
+ }
+
+ @Test
+ void testCreateNullVersionId() {
+ StorageLockData data = new StorageLockData(true, 1700000000000L, "testOwner");
+ IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () ->
+ new StorageLockFile(data, null)
+ );
+ assertTrue(exception.getMessage().contains("VersionId must not be null or empty."));
+ exception = assertThrows(IllegalArgumentException.class, () ->
+ new StorageLockFile(data, "")
+ );
+ assertTrue(exception.getMessage().contains("VersionId must not be null or empty."));
+ }
+
+ @Test
+ void testToJsonStreamValidData() {
+ StorageLockFile file = StorageLockFile.createFromStream(validJsonStream, VERSION_ID);
+ ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+ file.writeToStream(outputStream);
+ String outputJson = new String(outputStream.toByteArray());
+ assertTrue(outputJson.contains("\"expired\":false"));
+ assertTrue(outputJson.contains("\"validUntil\":1700000000000"));
+ assertTrue(outputJson.contains("\"owner\":\"testOwner\""));
+ }
+
+ @Test
+ void testToJsonStreamErrorHandling() throws IOException {
+ OutputStream mockOutputStream = mock(OutputStream.class);
+
+ doThrow(new IOException("Simulated IOException"))
+ .when(mockOutputStream)
+ .write(any(byte[].class), anyInt(), anyInt());
+ StorageLockFile file = new StorageLockFile(
+ new StorageLockData(true, System.currentTimeMillis() + 1000, "testOwner"),
+ VERSION_ID);
+
+ HoodieIOException exception = assertThrows(HoodieIOException.class, () -> file.writeToStream(mockOutputStream));
+ assertTrue(exception.getMessage().contains("Error writing object to JSON"));
+ }
+
+ @Test
+ void testToByteArrayValidData() {
+ StorageLockData data = new StorageLockData(false, 1700000000000L, "testOwner");
+ String outputJson = new String(StorageLockFile.toByteArray(data));
+ assertTrue(outputJson.contains("\"expired\":false"));
+ assertTrue(outputJson.contains("\"validUntil\":1700000000000"));
+ assertTrue(outputJson.contains("\"owner\":\"testOwner\""));
+ }
+
+ @Test
+ void testIsExpired() {
+ StorageLockData data = new StorageLockData(true, System.currentTimeMillis() - 1000, "testOwner");
+ StorageLockFile file = new StorageLockFile(data, VERSION_ID);
+ assertTrue(file.isExpired());
+ }
+
+ @Test
+ void testGetVersionId() {
+ StorageLockData data = new StorageLockData(false, 1700000000000L, "testOwner");
+ StorageLockFile file = new StorageLockFile(data, VERSION_ID);
+ assertEquals(VERSION_ID, file.getVersionId());
+ }
+
+ @Test
+ void testGetLockFolderPathWithoutTrailingSlash() {
+ String basePath = "s3://bucket/table";
+ String expected = "s3://bucket/table/.hoodie/.locks";
+ String actual = StorageLockClient.getLockFolderPath(basePath);
+ assertEquals(expected, actual);
+ }
+
+ @Test
+ void testGetLockFolderPathWithTrailingSlash() {
+ String basePath = "s3://bucket/table/";
+ String expected = "s3://bucket/table/.hoodie/.locks";
+ String actual = StorageLockClient.getLockFolderPath(basePath);
+ assertEquals(expected, actual, "Path with trailing slash should be normalized correctly");
+ }
+
+ @Test
+ void testGetLockFolderPathWithMultipleTrailingSlashes() {
+ String basePath = "s3://bucket/table///";
+ String expected = "s3://bucket/table/.hoodie/.locks";
+ String actual = StorageLockClient.getLockFolderPath(basePath);
+ assertEquals(expected, actual);
+ }
+
+ @Test
+ void testGetLockFolderPathLocalFileSystem() {
+ String basePath = "/tmp/hudi/table";
+ String expected = "/tmp/hudi/table/.hoodie/.locks";
+ String actual = StorageLockClient.getLockFolderPath(basePath);
+ assertEquals(expected, actual);
+ }
+
+ @Test
+ void testGetLockFolderPathLocalFileSystemWithTrailingSlash() {
+ String basePath = "/tmp/hudi/table/";
+ String expected = "/tmp/hudi/table/.hoodie/.locks";
+ String actual = StorageLockClient.getLockFolderPath(basePath);
+ assertEquals(expected, actual);
+ }
+
+ @Test
+ void testGetAuditConfigPathWithTrailingSlash() {
+ String basePath = "s3://bucket/table/";
+ String expected = "s3://bucket/table/.hoodie/.locks/audit_enabled.json";
+ String actual = StorageLockProviderAuditService.getAuditConfigPath(basePath);
+ assertEquals(expected, actual, "Audit config path with trailing slash should be normalized correctly");
+ }
+
+ @Test
+ void testGetAuditFolderPathWithTrailingSlash() {
+ String basePath = "s3://bucket/table/";
+ String expected = "s3://bucket/table/.hoodie/.locks/audit";
+ String actual = StorageLockProviderAuditService.getAuditFolderPath(basePath);
+ assertEquals(expected, actual, "Audit folder path with trailing slash should be normalized correctly");
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/models/TestLockProviderHeartbeatManager.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/models/TestLockProviderHeartbeatManager.java
new file mode 100644
index 000000000000..234aebe70b76
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/client/transaction/lock/models/TestLockProviderHeartbeatManager.java
@@ -0,0 +1,481 @@
+/*
+ * 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.hudi.client.transaction.lock.models;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+
+import static org.apache.hudi.client.transaction.lock.models.LockProviderHeartbeatManager.DEFAULT_STOP_HEARTBEAT_TIMEOUT_MS;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class TestLockProviderHeartbeatManager {
+
+ private ScheduledExecutorService mockScheduler;
+ private Logger mockLogger;
+ private ScheduledFuture<?> mockFuture;
+ private HeartbeatManager manager;
+ private static final String LOGGER_ID = "test-owner";
+ private ScheduledExecutorService actualExecutorService;
+
+ @BeforeEach
+ void setUp() {
+ mockScheduler = mock(ScheduledExecutorService.class);
+ mockLogger = mock(Logger.class);
+ mockFuture = mock(ScheduledFuture.class);
+ actualExecutorService = Executors.newSingleThreadScheduledExecutor(r -> {
+ Thread t = new Thread(r, "Heartbeat-Test-Thread");
+ t.setDaemon(true);
+ return t;
+ });
+ }
+
+ @AfterEach
+ void tearDown() throws Exception {
+ if (manager != null) {
+ manager.close();
+ manager = null;
+ }
+ actualExecutorService.shutdownNow();
+ }
+
+ @Test
+ void testStartHeartbeatSuccess() {
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), eq(100L), eq(100L), eq(TimeUnit.MILLISECONDS)))
+ .thenAnswer(invocation -> mockFuture);
+ manager = createDefaultManagerWithMocks(() -> true);
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ }
+
+ @Test
+ void testStartHeartbeatAlreadyRunning() {
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)))
+ .thenAnswer(invocation -> mockFuture);
+
+ manager = createDefaultManagerWithMocks(() -> true);
+
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ assertFalse(manager.startHeartbeatForThread(Thread.currentThread()));
+ verify(mockLogger).warn("Owner {}: Heartbeat is already running.", LOGGER_ID);
+ }
+
+ @Test
+ void testStartHeartbeatSchedulerException() {
+ doThrow(new RejectedExecutionException("Scheduler failure"))
+ .when(mockScheduler)
+ .scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class));
+
+ manager = createDefaultManagerWithMocks(() -> true);
+
+ assertFalse(manager.startHeartbeatForThread(Thread.currentThread()));
+ verify(mockLogger).error(eq("Owner {}: Unable to schedule heartbeat task. {}"), eq(LOGGER_ID), any(RejectedExecutionException.class));
+ }
+
+ @Test
+ void testStopHeartbeatNeverStarted() {
+ manager = createDefaultManagerWithMocks(() -> true);
+
+ assertFalse(manager.stopHeartbeat(true));
+ verify(mockLogger).warn("Owner {}: No active heartbeat task to stop.", LOGGER_ID);
+ }
+
+ @Test
+ void testStopHeartbeatAlreadyRequested() {
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)))
+ .thenAnswer(invocation -> mockFuture);
+
+ manager = createDefaultManagerWithMocks(() -> true);
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+
+ when(mockFuture.cancel(true)).thenReturn(true);
+ when(mockFuture.isDone()).thenReturn(false).thenReturn(true);
+
+ assertTrue(manager.stopHeartbeat(true));
+
+ // Call stop again
+ assertFalse(manager.stopHeartbeat(true));
+ verify(mockLogger).warn("Owner {}: No active heartbeat task to stop.", LOGGER_ID);
+ }
+
+ @Test
+ void testHeartbeatUnableToAcquireSemaphore() throws InterruptedException {
+ CountDownLatch latch = new CountDownLatch(1);
+ AtomicReference<Thread> t = new AtomicReference<>();
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)))
+ .thenAnswer(invocation -> {
+ Runnable task = invocation.getArgument(0);
+ t.set(new Thread(() -> {
+ task.run();
+ latch.countDown();
+ }));
+ return mockFuture;
+ });
+
+ when(mockFuture.cancel(true)).thenReturn(true);
+ Semaphore semaphore = mock(Semaphore.class);
+
+ // Stub the tryAcquire() method to return false (for the heartbeat) and true (for stop)
+ when(semaphore.tryAcquire()).thenReturn(false);
+ when(semaphore.tryAcquire(eq(DEFAULT_STOP_HEARTBEAT_TIMEOUT_MS), eq(TimeUnit.MILLISECONDS))).thenReturn(true);
+ manager = new LockProviderHeartbeatManager(
+ LOGGER_ID,
+ mockScheduler,
+ 100L,
+ DEFAULT_STOP_HEARTBEAT_TIMEOUT_MS,
+ () -> true,
+ semaphore,
+ mockLogger);
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ t.get().start();
+ assertTrue(latch.await(1000, TimeUnit.MILLISECONDS));
+ assertTrue(manager.stopHeartbeat(true));
+
+ verify(mockLogger).error("Owner {}: Heartbeat semaphore should be acquirable at the start of every heartbeat!", LOGGER_ID);
+ assertFalse(manager.hasActiveHeartbeat());
+ }
+
+ @Test
+ void testStopHeartbeatMockSuccessfulCancel() {
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)))
+ .thenAnswer(invocation -> mockFuture);
+ when(mockFuture.cancel(true)).thenReturn(true);
+
+ manager = createDefaultManagerWithMocks(() -> true);
+ manager.startHeartbeatForThread(Thread.currentThread());
+
+ when(mockFuture.isDone()).thenReturn(false).thenReturn(true);
+ assertTrue(manager.stopHeartbeat(true));
+ }
+
+ @Test
+ void testHeartbeatTaskHandlesInterrupt() throws InterruptedException {
+ CountDownLatch latch = new CountDownLatch(1);
+ AtomicReference<Thread> t = new AtomicReference<>();
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)))
+ .thenAnswer(invocation -> {
+ Runnable task = invocation.getArgument(0);
+ t.set(new Thread(() -> {
+ task.run();
+ latch.countDown();
+ }));
+ return mockFuture;
+ });
+
+ when(mockFuture.cancel(true)).thenReturn(true);
+
+ // Initialize heartbeat manager with a function that always returns false (renewal failure)
+ manager = createDefaultManagerWithMocks(() -> false);
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ t.get().start();
+ t.get().interrupt();
+
+ assertTrue(latch.await(500, TimeUnit.MILLISECONDS), "Heartbeat task did not run in time");
+
+ // This call will wait for heartbeat task to stop itself, as the semaphore has already been acquired by the heartbeat task.
+ assertFalse(manager.stopHeartbeat(true));
+
+ verify(mockLogger).warn("Owner {}: No active heartbeat task to stop.", LOGGER_ID);
+ verify(mockLogger).debug(
+ "Owner {}: Heartbeat started with interval: {} ms",
+ "test-owner",
+ 100L
+ );
+ verify(mockLogger).info("Owner {}: Requested termination of heartbeat task. Cancellation returned {}.", LOGGER_ID, true);
+ assertFalse(manager.hasActiveHeartbeat());
+ }
+
+ @Test
+ void testHeartbeatTaskNullWriter() {
+ manager = createDefaultManagerWithMocks(() -> true);
+ assertThrows(IllegalArgumentException.class, () -> manager.startHeartbeatForThread(null));
+ }
+
+ @Test
+ void testHeartbeatTaskImmediateDeadMonitoringThread() throws InterruptedException {
+ // Use a real thread that will terminate immediately.
+ CountDownLatch latch = new CountDownLatch(1);
+ AtomicReference<Thread> t = new AtomicReference<>();
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)))
+ .thenAnswer(invocation -> {
+ Runnable task = invocation.getArgument(0);
+ t.set(new Thread(() -> {
+ task.run();
+ latch.countDown();
+ }));
+ return mockFuture;
+ });
+
+ when(mockFuture.cancel(false)).thenReturn(false);
+ Thread deadThread = new Thread(() -> {
+ });
+ deadThread.start();
+ deadThread.join();
+ manager = createDefaultManagerWithMocks(() -> true);
+
+ assertTrue(manager.startHeartbeatForThread(deadThread));
+ t.get().start();
+
+ assertTrue(latch.await(500, TimeUnit.MILLISECONDS), "Heartbeat task did not run in time");
+ verify(mockLogger).warn("Owner {}: Monitored thread is no longer alive.", LOGGER_ID);
+ verify(mockLogger).info("Owner {}: Requested termination of heartbeat task. Cancellation returned {}.", LOGGER_ID, false);
+ assertFalse(manager.hasActiveHeartbeat());
+ }
+
+ @Test
+ void testHeartbeatTaskRenewalException() throws InterruptedException {
+ CountDownLatch latch = new CountDownLatch(1);
+ AtomicReference<Thread> t = new AtomicReference<>();
+ when(mockScheduler.scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)))
+ .thenAnswer(invocation -> {
+ Runnable task = invocation.getArgument(0);
+ t.set(new Thread(() -> {
+ task.run();
+ latch.countDown();
+ }));
+ return mockFuture;
+ });
+ manager = createDefaultManagerWithMocks(() -> {
+ throw new RuntimeException("Renewal error");
+ });
+
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ t.get().start();
+ assertTrue(latch.await(500, TimeUnit.MILLISECONDS), "Heartbeat task did not run in time");
+ verify(mockLogger).error(
+ eq("Owner {}: Heartbeat function threw exception {}"),
+ eq(LOGGER_ID),
+ any(RuntimeException.class));
+ assertFalse(manager.hasActiveHeartbeat());
+ }
+
+ @Test
+ void testHeartbeatStopWaitsForHeartbeatTaskToFinish() throws InterruptedException {
+ // Use a real thread
+ CountDownLatch stopHeartbeatTaskLatch = new CountDownLatch(1);
+ manager = createDefaultManagerWithRealExecutor(() -> {
+ try {
+ // This will freeze the heartbeat task.
+ assertTrue(stopHeartbeatTaskLatch.await(500, TimeUnit.MILLISECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return true;
+ });
+
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ CountDownLatch finishStopHeartbeatLatch = new CountDownLatch(1);
+ Thread t = new Thread(() -> {
+ assertTrue(manager.stopHeartbeat(false));
+ finishStopHeartbeatLatch.countDown();
+ });
+ t.start();
+ // Unblock the heartbeat task.
+ stopHeartbeatTaskLatch.countDown();
+ assertTrue(finishStopHeartbeatLatch.await(500, TimeUnit.MILLISECONDS), "Stop heartbeat task did not finish.");
+ assertFalse(manager.hasActiveHeartbeat());
+ verify(mockLogger).debug("Owner {}: Heartbeat task successfully terminated.", LOGGER_ID);
+ }
+
+ @Test
+ void testHeartbeatUnableToStopHeartbeatTask() throws InterruptedException {
+ CountDownLatch stopHeartbeatTaskLatch = new CountDownLatch(1);
+ CountDownLatch heartbeatStartedLatch = new CountDownLatch(1);
+ // Set stop heartbeat timeout to 5000ms
+ manager = new LockProviderHeartbeatManager(LOGGER_ID, actualExecutorService, 100L, 5000L, () -> {
+ try {
+ // Tells us that the heartbeat has started
+ heartbeatStartedLatch.countDown();
+ // This will freeze the heartbeat task.
+ assertTrue(stopHeartbeatTaskLatch.await(10000, TimeUnit.MILLISECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ // Regardless of whether we return true or false the future executions will be cancelled.
+ return true;
+ }, new Semaphore(1), mockLogger);
+
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ CountDownLatch stopHeartbeatLatch = new CountDownLatch(1);
+ Thread stopHeartbeatThread = new Thread(() -> {
+ // Try to stop the heartbeat (this should hang for 15 seconds)
+ assertFalse(manager.stopHeartbeat(false));
+ stopHeartbeatLatch.countDown();
+ });
+ assertTrue(heartbeatStartedLatch.await(500, TimeUnit.MILLISECONDS), "Heartbeat task did not start.");
+ stopHeartbeatThread.start();
+ assertTrue(stopHeartbeatLatch.await(7000, TimeUnit.MILLISECONDS), "Stop heartbeat task did not finish.");
+ assertTrue(manager.hasActiveHeartbeat());
+ verify(mockLogger).error("Owner {}: Heartbeat is still in flight!", LOGGER_ID);
+ // Unblock the heartbeat task.
+ stopHeartbeatTaskLatch.countDown();
+ }
+
+ @Test
+ void testHeartbeatInterruptStopHeartbeatTask() throws InterruptedException {
+ CountDownLatch stopHeartbeatTaskLatch = new CountDownLatch(1);
+ CountDownLatch heartbeatStartedLatch = new CountDownLatch(1);
+ // Set stop heartbeat timeout to 5000ms
+ manager = new LockProviderHeartbeatManager(LOGGER_ID, actualExecutorService, 100L, 5000L, () -> {
+ try {
+ // Tells us that the heartbeat has started
+ heartbeatStartedLatch.countDown();
+ // This will freeze the heartbeat task.
+ assertTrue(stopHeartbeatTaskLatch.await(10000, TimeUnit.MILLISECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ // Regardless of whether we return true or false the future executions will be cancelled.
+ return true;
+ }, new Semaphore(1), mockLogger);
+
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ CountDownLatch stopHeartbeatLatch = new CountDownLatch(1);
+ Thread stopHeartbeatThread = new Thread(() -> {
+ // Try to stop the heartbeat (this should hang for 15 seconds)
+ assertFalse(manager.stopHeartbeat(false));
+ stopHeartbeatLatch.countDown();
+ });
+ assertTrue(heartbeatStartedLatch.await(500, TimeUnit.MILLISECONDS), "Heartbeat task did not start.");
+ stopHeartbeatThread.start();
+ stopHeartbeatThread.interrupt();
+ assertTrue(stopHeartbeatLatch.await(7000, TimeUnit.MILLISECONDS), "Stop heartbeat task did not finish.");
+ assertTrue(manager.hasActiveHeartbeat());
+ verify(mockLogger).warn("Owner {}: Interrupted while waiting for heartbeat termination.", LOGGER_ID);
+ // Unblock the heartbeat task.
+ stopHeartbeatTaskLatch.countDown();
+ }
+
+ @Test
+ void testHeartbeatTaskValidateStop() throws InterruptedException {
+ CountDownLatch latch = new CountDownLatch(2);
+
+ manager = createDefaultManagerWithRealExecutor(() -> {
+ latch.countDown();
+ return true;
+ });
+
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+
+ // Wait until at least two heartbeat renewals have occurred
+ assertTrue(latch.await(2000, TimeUnit.MILLISECONDS), "Heartbeat did not renew twice in time");
+
+ assertEquals(0, latch.getCount(), "Heartbeat did not execute exactly twice");
+
+ assertTrue(manager.hasActiveHeartbeat());
+ assertTrue(manager.stopHeartbeat(false));
+ assertFalse(manager.hasActiveHeartbeat());
+ }
+
+ @Test
+ void testDefaultManagerRapidStartStop1Ms() {
+ manager = new LockProviderHeartbeatManager(LOGGER_ID, 1, () -> true);
+
+ for (int i = 0; i < 100; i++) {
+ assertTrue(manager.startHeartbeatForThread(Thread.currentThread()));
+ assertTrue(manager.hasActiveHeartbeat());
+ assertTrue(manager.stopHeartbeat(true));
+ assertFalse(manager.hasActiveHeartbeat());
+ }
+ }
+
+ @Test
+ void testClose() throws Exception {
+ manager = createDefaultManagerWithMocks(() -> true);
+ manager.close();
+ assertFalse(manager.hasActiveHeartbeat());
+ }
+
+ @Test
+ void testClose_StopsHeartbeatAndShutsDownScheduler() throws Exception {
+ when(mockScheduler.awaitTermination(5, TimeUnit.SECONDS)).thenReturn(true);
+ manager = createDefaultManagerWithMocks(() -> true);
+
+ manager.close();
+
+ verify(mockScheduler).shutdown();
+ verify(mockScheduler, never()).shutdownNow();
+ }
+
+ @Test
+ void testClose_ForceShutdownWhenTerminationTimesOut() throws Exception {
+ when(mockScheduler.awaitTermination(5, TimeUnit.SECONDS)).thenReturn(false);
+ manager = createDefaultManagerWithMocks(() -> true);
+
+ manager.close();
+
+ verify(mockScheduler).shutdown();
+ verify(mockScheduler).shutdownNow();
+ }
+
+ @Test
+ void testClose_HandlesInterruptedException() throws Exception {
+ when(mockScheduler.awaitTermination(5, TimeUnit.SECONDS)).thenThrow(new InterruptedException());
+ manager = createDefaultManagerWithMocks(() -> true);
+
+ manager.close();
+
+ verify(mockScheduler).shutdown();
+ verify(mockScheduler).shutdownNow();
+ assertTrue(Thread.currentThread().isInterrupted(), "Thread should be interrupted after exception handling");
+ }
+
+ private LockProviderHeartbeatManager createDefaultManagerWithMocks(Supplier<Boolean> heartbeatFunc) {
+ return new LockProviderHeartbeatManager(
+ LOGGER_ID,
+ mockScheduler,
+ 100L,
+ DEFAULT_STOP_HEARTBEAT_TIMEOUT_MS,
+ heartbeatFunc,
+ new Semaphore(1),
+ mockLogger);
+ }
+
+ private LockProviderHeartbeatManager createDefaultManagerWithRealExecutor(Supplier<Boolean> heartbeatFunc) {
+ return new LockProviderHeartbeatManager(
+ LOGGER_ID,
+ actualExecutorService,
+ 100L,
+ DEFAULT_STOP_HEARTBEAT_TIMEOUT_MS,
+ heartbeatFunc,
+ new Semaphore(1),
+ mockLogger);
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/config/TestStorageBasedLockConfig.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/config/TestStorageBasedLockConfig.java
new file mode 100644
index 000000000000..1ccf6e665942
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/config/TestStorageBasedLockConfig.java
@@ -0,0 +1,101 @@
+/*
+ * 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.hudi.config;
+
+import org.apache.hudi.common.config.TypedProperties;
+
+import org.junit.jupiter.api.Test;
+
+import static org.apache.hudi.config.HoodieWriteConfig.BASE_PATH;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class TestStorageBasedLockConfig {
+
+ @Test
+ void testDefaultValues() {
+ // Asserts that the defaults are correct
+ TypedProperties props = new TypedProperties();
+ props.setProperty(BASE_PATH.key(), "s3://hudi-bucket/table/basepath");
+
+ StorageBasedLockConfig.Builder builder = new StorageBasedLockConfig.Builder();
+ StorageBasedLockConfig config = builder
+ .fromProperties(props)
+ .build();
+
+ assertEquals(5 * 60, config.getValiditySeconds(), "Default lock validity should be 5 minutes");
+ assertEquals(30, config.getRenewIntervalSecs(), "Default heartbeat poll time should be 30 seconds");
+ }
+
+ @Test
+ void testCustomValues() {
+ // Testing that custom values which differ from defaults can be read properly
+ TypedProperties props = new TypedProperties();
+ props.setProperty(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "120");
+ props.setProperty(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "10");
+ props.setProperty(BASE_PATH.key(), "/hudi/table/basepath");
+
+ StorageBasedLockConfig config = new StorageBasedLockConfig.Builder()
+ .fromProperties(props)
+ .build();
+
+ assertEquals(120, config.getValiditySeconds());
+ assertEquals(10, config.getRenewIntervalSecs());
+ assertEquals("/hudi/table/basepath", config.getHudiTableBasePath());
+ }
+
+ @Test
+ void testBasePathPropertiesValidation() {
+ // Tests that validations around the base path are present.
+ TypedProperties props = new TypedProperties();
+ props.setProperty(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "120");
+ StorageBasedLockConfig.Builder propsBuilder = new StorageBasedLockConfig.Builder();
+
+ // Missing base path
+ IllegalArgumentException exception = assertThrows(IllegalArgumentException.class,
+ () -> propsBuilder.fromProperties(props));
+ assertTrue(exception.getMessage().contains(BASE_PATH.key()));
+ }
+
+ @Test
+ void testTimeThresholds() {
+ // Ensure that validations which restrict the time-based inputs are working.
+ TypedProperties props = new TypedProperties();
+ props.setProperty(BASE_PATH.key(), "/hudi/table/basepath");
+ // Invalid config case: validity timeout is less than 10x of heartbeat poll period
+ props.setProperty(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "5");
+ props.setProperty(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "3");
+ StorageBasedLockConfig.Builder propsBuilder = new StorageBasedLockConfig.Builder();
+
+ IllegalArgumentException exception = assertThrows(IllegalArgumentException.class,
+ () -> propsBuilder.fromProperties(props));
+ assertTrue(exception.getMessage().contains(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key()));
+ // Invalid config case: validity timeout is less than 10 seconds
+ props.setProperty(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "9");
+ props.setProperty(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "1");
+ exception = assertThrows(IllegalArgumentException.class, () -> propsBuilder.fromProperties(props));
+ assertTrue(exception.getMessage().contains(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key()));
+ // Invalid config case: heartbeat poll period is less than 1 second
+ props.setProperty(StorageBasedLockConfig.VALIDITY_TIMEOUT_SECONDS.key(), "10");
+ props.setProperty(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key(), "0");
+ exception = assertThrows(IllegalArgumentException.class, () -> propsBuilder.fromProperties(props));
+ assertTrue(exception.getMessage().contains(StorageBasedLockConfig.RENEW_INTERVAL_SECS.key()));
+ }
+}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestMetricsReporterFactory.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestMetricsReporterFactory.java
index a44443d9bd5d..21d8a0901617 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestMetricsReporterFactory.java
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestMetricsReporterFactory.java
@@ -96,4 +96,3 @@ public class TestMetricsReporterFactory {
}
}
-
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/fs/StorageSchemes.java b/hudi-common/src/main/java/org/apache/hudi/common/fs/StorageSchemes.java
index 54a9853d5169..c17198a1e5d7 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/fs/StorageSchemes.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/fs/StorageSchemes.java
@@ -20,64 +20,67 @@ package org.apache.hudi.common.fs;
import java.util.Arrays;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.StringUtils;
+
/**
* All the supported storage schemes in Hoodie.
*/
public enum StorageSchemes {
// Local filesystem
- FILE("file", false, false, true, true),
+ FILE("file", false, false, true, true, null),
// Hadoop File System
- HDFS("hdfs", true, false, true, false),
+ HDFS("hdfs", true, false, true, false, null),
// Baidu Advanced File System
- AFS("afs", true, null, null, null),
+ AFS("afs", true, null, null, null, null),
// Mapr File System
- MAPRFS("maprfs", true, null, null, null),
+ MAPRFS("maprfs", true, null, null, null, null),
// Apache Ignite FS
- IGNITE("igfs", true, null, null, null),
+ IGNITE("igfs", true, null, null, null, null),
// AWS S3
- S3A("s3a", false, true, null, true),
- S3("s3", false, true, null, true),
+ S3A("s3a", false, true, null, true, "org.apache.hudi.aws.transaction.lock.S3StorageLockClient"),
+ S3("s3", false, true, null, true, "org.apache.hudi.aws.transaction.lock.S3StorageLockClient"),
// Google Cloud Storage
- GCS("gs", false, true, null, true),
+ GCS("gs", false, true, null, true, "org.apache.hudi.gcp.transaction.lock.GCSStorageLockClient"),
// Azure WASB
- WASB("wasb", false, null, null, null), WASBS("wasbs", false, null, null, null),
+ WASB("wasb", false, null, null, null, null), WASBS("wasbs", false, null, null, null, null),
// Azure ADLS
- ADL("adl", false, null, null, null),
+ ADL("adl", false, null, null, null, null),
// Azure ADLS Gen2
- ABFS("abfs", false, null, null, null), ABFSS("abfss", false, null, null, null),
+ ABFS("abfs", false, null, null, null, null), ABFSS("abfss", false, null, null, null, null),
// Aliyun OSS
- OSS("oss", false, null, null, null),
+ OSS("oss", false, null, null, null, null),
// View FS for federated setups. If federating across cloud stores, then append support is false
// View FS support atomic creation
- VIEWFS("viewfs", true, null, true, null),
+ VIEWFS("viewfs", true, null, true, null, null),
//ALLUXIO
- ALLUXIO("alluxio", false, null, null, null),
+ ALLUXIO("alluxio", false, null, null, null, null),
// Tencent Cloud Object Storage
- COSN("cosn", false, null, null, null),
+ COSN("cosn", false, null, null, null, null),
// Tencent Cloud HDFS
- CHDFS("ofs", true, null, null, null),
+ CHDFS("ofs", true, null, null, null, null),
// Tencent Cloud CacheFileSystem
- GOOSEFS("gfs", false, null, null, null),
+ GOOSEFS("gfs", false, null, null, null, null),
// Databricks file system
- DBFS("dbfs", false, null, null, null),
+ DBFS("dbfs", false, null, null, null, null),
// IBM Cloud Object Storage
- COS("cos", false, null, null, null),
+ COS("cos", false, null, null, null, null),
// Huawei Cloud Object Storage
- OBS("obs", false, null, null, null),
+ OBS("obs", false, null, null, null, null),
// Kingsoft Standard Storage ks3
- KS3("ks3", false, null, null, null),
+ KS3("ks3", false, null, null, null, null),
// JuiceFileSystem
- JFS("jfs", true, null, null, null),
+ JFS("jfs", true, null, null, null, null),
// Baidu Object Storage
- BOS("bos", false, null, null, null),
+ BOS("bos", false, null, null, null, null),
// Oracle Cloud Infrastructure Object Storage
- OCI("oci", false, null, null, null),
+ OCI("oci", false, null, null, null, null),
// Volcengine Object Storage
- TOS("tos", false, null, null, null),
+ TOS("tos", false, null, null, null, null),
// Volcengine Cloud HDFS
- CFS("cfs", true, null, null, null),
+ CFS("cfs", true, null, null, null, null),
// Aliyun Apsara File Storage for HDFS
- DFS("dfs", true, false, true, null);
+ DFS("dfs", true, false, true, null, null);
private String scheme;
private boolean supportsAppend;
@@ -89,13 +92,15 @@ public enum StorageSchemes {
// when we want to get only part of files under a directory rather than all files, use getStatus may be more friendly than listStatus.
// here is a trade-off between rpc times and throughput of storage meta service
private Boolean listStatusFriendly;
+ private String storageLockClass;
- StorageSchemes(String scheme, boolean supportsAppend, Boolean isWriteTransactional, Boolean supportAtomicCreation, Boolean listStatusFriendly) {
+ StorageSchemes(String scheme, boolean supportsAppend, Boolean isWriteTransactional, Boolean supportAtomicCreation, Boolean listStatusFriendly, String storageLockClass) {
this.scheme = scheme;
this.supportsAppend = supportsAppend;
this.isWriteTransactional = isWriteTransactional;
this.supportAtomicCreation = supportAtomicCreation;
this.listStatusFriendly = listStatusFriendly;
+ this.storageLockClass = storageLockClass;
}
public String getScheme() {
@@ -118,6 +123,14 @@ public enum StorageSchemes {
return listStatusFriendly != null && listStatusFriendly;
}
+ public boolean implementsStorageLock() {
+ return !StringUtils.isNullOrEmpty(storageLockClass);
+ }
+
+ public String getStorageLockClass() {
+ return storageLockClass;
+ }
+
public static boolean isSchemeSupported(String scheme) {
return Arrays.stream(values()).anyMatch(s -> s.getScheme().equals(scheme));
}
@@ -150,4 +163,12 @@ public enum StorageSchemes {
}
return Arrays.stream(StorageSchemes.values()).anyMatch(s -> s.getListStatusFriendly() && s.scheme.equals(scheme));
}
+
+ public static Option<StorageSchemes> getStorageLockImplementationIfExists(String scheme) {
+ if (!isSchemeSupported(scheme)) {
+ throw new IllegalArgumentException("Unsupported scheme :" + scheme);
+ }
+ return Option.fromJavaOptional(Arrays.stream(StorageSchemes.values())
+ .filter(s -> s.implementsStorageLock() && s.scheme.equals(scheme)).findFirst());
+ }
}
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java b/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java
index c84ce4a6320d..f8a6b2d64afa 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java
@@ -99,6 +99,8 @@ public class HoodieTableMetaClient implements Serializable {
public static final String BOOTSTRAP_INDEX_BY_FILE_ID_FOLDER_PATH = BOOTSTRAP_INDEX_ROOT_FOLDER_PATH + Path.SEPARATOR
+ ".fileids";
+ public static final String LOCKS_FOLDER_NAME = METAFOLDER_NAME + Path.SEPARATOR + ".locks";
+
public static final String SCHEMA_FOLDER_NAME = ".schema";
public static final String MARKER_EXTN = ".marker";
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/util/HoodieTimer.java b/hudi-common/src/main/java/org/apache/hudi/common/util/HoodieTimer.java
index a35cd60933f8..42ad16065f42 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/util/HoodieTimer.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/util/HoodieTimer.java
@@ -83,6 +83,17 @@ public class HoodieTimer {
return timeInfoDeque.pop().stop();
}
+ /**
+ * tries to end the timer but will not cause a runtime failure like endTimer will if the queue is empty
+ * @return option of the timer duration, empty option if the queue is empty
+ */
+ public Option<Long> tryEndTimer() {
+ if (timeInfoDeque.isEmpty()) {
+ return Option.empty();
+ }
+ return Option.of(timeInfoDeque.pop().stop());
+ }
+
/**
* Creates an instance of {@link HoodieTimer} already started
*/
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/util/Option.java b/hudi-common/src/main/java/org/apache/hudi/common/util/Option.java
index 957dab28e2c2..619d389d884d 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/util/Option.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/util/Option.java
@@ -84,6 +84,10 @@ public final class Option<T> implements Serializable {
return null != val;
}
+ public boolean isEmpty() {
+ return !isPresent();
+ }
+
public T get() {
if (null == val) {
throw new NoSuchElementException("No value present in Option");
diff --git a/hudi-common/src/test/java/org/apache/hudi/common/fs/TestFSUtils.java b/hudi-common/src/test/java/org/apache/hudi/common/fs/TestFSUtils.java
index 7ad2bb3cb900..587fead6e1a0 100644
--- a/hudi-common/src/test/java/org/apache/hudi/common/fs/TestFSUtils.java
+++ b/hudi-common/src/test/java/org/apache/hudi/common/fs/TestFSUtils.java
@@ -41,6 +41,8 @@ import org.junit.contrib.java.lang.system.EnvironmentVariables;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import java.io.IOException;
import java.nio.file.Files;
@@ -619,6 +621,24 @@ public class TestFSUtils extends HoodieCommonTestHarness {
assertEquals("s3://my-bucket/s3a://another-bucket/another/path", FSUtils.s3aToS3("s3a://my-bucket/s3a://another-bucket/another/path"));
}
+ @ParameterizedTest
+ @ValueSource(strings = {
+ "gs://my-bucket/path/to/object",
+ "gs://my-bucket",
+ "gs://MY-BUCKET/PATH/TO/OBJECT",
+ "https://myaccount.blob.core.windows.net/mycontainer/path/to/blob",
+ "https://myaccount.blob.core.windows.net/MYCONTAINER/PATH/TO/BLOB",
+ "https://example.com/path/to/resource",
+ "http://example.com",
+ "ftp://example.com/resource",
+ "",
+ "gs://my-bucket/path/to/s3a://object",
+ "gs://my-bucket s3a://my-object",
+ })
+ void testUriDoesNotChange(String uri) {
+ assertEquals(uri, FSUtils.s3aToS3(uri));
+ }
+
@Test
void testGetPathWithoutScheme() {
String path1 = "s3://test_bucket_one/table/base/path";
diff --git a/hudi-common/src/test/java/org/apache/hudi/common/fs/TestStorageSchemes.java b/hudi-common/src/test/java/org/apache/hudi/common/fs/TestStorageSchemes.java
index 7ac8a9bcabb6..8336d051d3f5 100644
--- a/hudi-common/src/test/java/org/apache/hudi/common/fs/TestStorageSchemes.java
+++ b/hudi-common/src/test/java/org/apache/hudi/common/fs/TestStorageSchemes.java
@@ -79,6 +79,15 @@ public class TestStorageSchemes {
assertThrows(IllegalArgumentException.class, () -> {
StorageSchemes.isAppendSupported("s2");
}, "Should throw exception for unsupported schemes");
+
+ for (StorageSchemes scheme : StorageSchemes.values()) {
+ String schemeName = scheme.getScheme();
+ if (scheme.getScheme().startsWith("s3") || scheme.getScheme().startsWith("gs")) {
+ assertTrue(StorageSchemes.getStorageLockImplementationIfExists(schemeName).isPresent());
+ } else {
+ assertFalse(StorageSchemes.getStorageLockImplementationIfExists(schemeName).isPresent());
+ }
+ }
}
@Test
diff --git a/hudi-aws/src/main/java/org/apache/hudi/aws/utils/S3Utils.java b/hudi-common/src/test/java/org/apache/hudi/common/util/TestHoodieTimer.java
similarity index 58%
rename from hudi-aws/src/main/java/org/apache/hudi/aws/utils/S3Utils.java
rename to hudi-common/src/test/java/org/apache/hudi/common/util/TestHoodieTimer.java
index bfb208ee1505..bc3c7d1446c6 100644
--- a/hudi-aws/src/main/java/org/apache/hudi/aws/utils/S3Utils.java
+++ b/hudi-common/src/test/java/org/apache/hudi/common/util/TestHoodieTimer.java
@@ -17,11 +17,25 @@
* under the License.
*/
-package org.apache.hudi.aws.utils;
+package org.apache.hudi.common.util;
-public final class S3Utils {
+import org.junit.jupiter.api.Test;
- public static String s3aToS3(String s3aUrl) {
- return s3aUrl.replaceFirst("(?i)^s3a://", "s3://");
+import static java.lang.Thread.sleep;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class TestHoodieTimer {
+
+ @Test
+ public void testTryStop() throws InterruptedException {
+ HoodieTimer timer = new HoodieTimer();
+ assertFalse(timer.tryEndTimer().isPresent());
+
+ timer.startTimer();
+ sleep(250);
+ Option<Long> result = timer.tryEndTimer();
+ assertTrue(result.isPresent());
+ assertTrue(result.get() >= 250);
}
}
diff --git a/hudi-gcp/pom.xml b/hudi-gcp/pom.xml
index 190a74848082..5534af6ec0e3 100644
--- a/hudi-gcp/pom.xml
+++ b/hudi-gcp/pom.xml
@@ -61,6 +61,11 @@ See https://github.com/GoogleCloudPlatform/cloud-opensource-java/wiki/The-Google
<artifactId>hudi-sync-common</artifactId>
<version>${project.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.hudi</groupId>
+ <artifactId>hudi-client-common</artifactId>
+ <version>${project.version}</version>
+ </dependency>
<dependency>
<groupId>com.google.cloud</groupId>
@@ -77,6 +82,10 @@ See https://github.com/GoogleCloudPlatform/cloud-opensource-java/wiki/The-Google
<artifactId>gcs-connector</artifactId>
<version>${gcs.connector.version}</version>
</dependency>
+ <dependency>
+ <groupId>com.google.cloud</groupId>
+ <artifactId>google-cloud-storage</artifactId>
+ </dependency>
<dependency>
<groupId>org.apache.parquet</groupId>
@@ -95,6 +104,14 @@ See https://github.com/GoogleCloudPlatform/cloud-opensource-java/wiki/The-Google
<artifactId>hadoop-common</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.hudi</groupId>
+ <artifactId>hudi-client-common</artifactId>
+ <version>${project.version}</version>
+ <classifier>tests</classifier>
+ <scope>test</scope>
+ </dependency>
+
<dependency>
<groupId>org.apache.hudi</groupId>
<artifactId>hudi-tests-common</artifactId>
diff --git a/hudi-gcp/src/main/java/org/apache/hudi/gcp/transaction/lock/GCSStorageLockClient.java b/hudi-gcp/src/main/java/org/apache/hudi/gcp/transaction/lock/GCSStorageLockClient.java
new file mode 100644
index 000000000000..1e782cc109dd
--- /dev/null
+++ b/hudi-gcp/src/main/java/org/apache/hudi/gcp/transaction/lock/GCSStorageLockClient.java
@@ -0,0 +1,289 @@
+/*
+ * 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.hudi.gcp.transaction.lock;
+
+import org.apache.hudi.client.transaction.lock.StorageLockClient;
+import org.apache.hudi.client.transaction.lock.models.LockGetResult;
+import org.apache.hudi.client.transaction.lock.models.LockUpsertResult;
+import org.apache.hudi.client.transaction.lock.models.StorageLockData;
+import org.apache.hudi.client.transaction.lock.models.StorageLockFile;
+import org.apache.hudi.common.util.Functions;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.VisibleForTesting;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.exception.HoodieIOException;
+
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.BlobId;
+import com.google.cloud.storage.BlobInfo;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
+import com.google.cloud.storage.StorageOptions;
+import org.jetbrains.annotations.NotNull;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
+import java.nio.channels.Channels;
+import java.util.Properties;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+/**
+ * A GCS-based implementation of a distributed lock provider using conditional writes
+ * with generationMatch, plus local concurrency safety, heartbeat/renew, and pruning old locks.
+ */
+@ThreadSafe
+public class GCSStorageLockClient implements StorageLockClient {
+ private static final Logger DEFAULT_LOGGER = LoggerFactory.getLogger(GCSStorageLockClient.class);
+ private static final long PRECONDITION_FAILURE_ERROR_CODE = 412;
+ private static final long NOT_FOUND_ERROR_CODE = 404;
+ private static final long RATE_LIMIT_ERROR_CODE = 429;
+ private static final long INTERNAL_SERVER_ERROR_CODE_MIN = 500;
+ private final Logger logger;
+ private final Storage gcsClient;
+ private final String bucketName;
+ private final String lockFilePath;
+ private final String ownerId;
+
+ /**
+ * Constructor that is used by reflection to instantiate a GCS-based locking service.
+ *
+ * @param ownerId The owner id.
+ * @param lockFileUri The path within the bucket where to write lock files.
+ * @param props The properties for the lock config, can be used to customize client.
+ */
+ public GCSStorageLockClient(
+ String ownerId,
+ String lockFileUri,
+ Properties props) {
+ this(ownerId, lockFileUri, props, createDefaultGcsClient(), DEFAULT_LOGGER);
+ }
+
+ @VisibleForTesting
+ GCSStorageLockClient(
+ String ownerId,
+ String lockFileUri,
+ Properties properties,
+ Functions.Function1<Properties, Storage> gcsClientSupplier,
+ Logger logger) {
+ Pair<String, String> bucketAndPath = StorageLockClient.parseBucketAndPath(lockFileUri);
+ this.bucketName = bucketAndPath.getLeft();
+ this.lockFilePath = bucketAndPath.getRight();
+ this.gcsClient = gcsClientSupplier.apply(properties);
+ this.ownerId = ownerId;
+ this.logger = logger;
+ }
+
+ private static Functions.Function1<Properties, Storage> createDefaultGcsClient() {
+ return (props) -> {
+ // Provide the option to customize the timeouts later on.
+ // For now, defaults suffice
+ return StorageOptions.newBuilder().build().getService();
+ };
+ }
+
+ /**
+ * Attempts to create or update the lock file using the given lock data and generation number.
+ *
+ * @param lockData the new lock data to use.
+ * @param generationNumber the expected generation number (0 for creation).
+ * @return the updated StorageLockFile instance.
+ * @throws StorageException if the update fails.
+ */
+ private StorageLockFile createOrUpdateLockFileInternal(StorageLockData lockData, long generationNumber)
+ throws StorageException {
+ BlobInfo blobInfo = BlobInfo.newBuilder(BlobId.of(bucketName, lockFilePath)).build();
+ Blob updatedBlob = gcsClient.create(
+ blobInfo,
+ StorageLockFile.toByteArray(lockData),
+ Storage.BlobTargetOption.generationMatch(generationNumber));
+ return new StorageLockFile(
+ lockData,
+ String.valueOf(updatedBlob.getGeneration()));
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Pair<LockUpsertResult, Option<StorageLockFile>> tryUpsertLockFile(
+ StorageLockData newLockData,
+ Option<StorageLockFile> previousLockFile) {
+ long generationNumber = getGenerationNumber(previousLockFile);
+ try {
+ StorageLockFile updatedFile = createOrUpdateLockFileInternal(newLockData, generationNumber);
+ return Pair.of(LockUpsertResult.SUCCESS, Option.of(updatedFile));
+ } catch (StorageException e) {
+ if (e.getCode() == PRECONDITION_FAILURE_ERROR_CODE) {
+ logger.info("OwnerId: {}, Unable to write new lock file. Another process has modified this lockfile {} already.",
+ ownerId, lockFilePath);
+ return Pair.of(LockUpsertResult.ACQUIRED_BY_OTHERS, Option.empty());
+ } else if (e.getCode() == RATE_LIMIT_ERROR_CODE) {
+ logger.warn("OwnerId: {}, Rate limit exceeded for lock file: {}", ownerId, lockFilePath);
+ } else if (e.getCode() >= INTERNAL_SERVER_ERROR_CODE_MIN) {
+ logger.warn("OwnerId: {}, GCS returned internal server error code for lock file: {}",
+ ownerId, lockFilePath, e);
+ } else {
+ throw e;
+ }
+ return Pair.of(LockUpsertResult.UNKNOWN_ERROR, Option.empty());
+ }
+ }
+
+ /**
+ * Handling storage exception for GET request
+ *
+ * @param e The error to handle.
+ * @param ignore404 Whether to ignore 404 as a valid exception.
+ * When we read from stream we might see this, and
+ * it should not be counted as NOT_EXISTS.
+ * @return The type of getResult error
+ */
+ private LockGetResult handleGetStorageException(StorageException e, boolean ignore404) {
+ if (e.getCode() == NOT_FOUND_ERROR_CODE) {
+ if (ignore404) {
+ logger.info("OwnerId: {}, GCS stream read failure detected: {}", ownerId, lockFilePath);
+ } else {
+ logger.info("OwnerId: {}, Object not found in the path: {}", ownerId, lockFilePath);
+ return LockGetResult.NOT_EXISTS;
+ }
+ } else if (e.getCode() == RATE_LIMIT_ERROR_CODE) {
+ logger.warn("OwnerId: {}, Rate limit exceeded for lock file: {}", ownerId, lockFilePath);
+ } else if (e.getCode() >= INTERNAL_SERVER_ERROR_CODE_MIN) {
+ logger.warn("OwnerId: {}, GCS returned internal server error code for lock file: {}", ownerId, lockFilePath, e);
+ } else {
+ throw e;
+ }
+ return LockGetResult.UNKNOWN_ERROR;
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Pair<LockGetResult, Option<StorageLockFile>> readCurrentLockFile() {
+ try {
+ Blob blob = gcsClient.get(BlobId.of(bucketName, lockFilePath));
+ if (blob == null) {
+ return Pair.of(LockGetResult.NOT_EXISTS, Option.empty());
+ }
+ return getLockFileFromBlob(blob);
+ } catch (StorageException e) {
+ return Pair.of(handleGetStorageException(e, false), Option.empty());
+ } catch (HoodieIOException e) {
+ // GCS will throw IOException wrapping 404 when reading from stream for file that has been modified in between calling gcsClient.get
+ // and StorageLockFile.createFromStream. People have complained that this is not being strongly consistent, however
+ // we have to handle this case. https://stackoverflow.com/q/66759993
+ Throwable cause = e.getCause();
+ if (cause instanceof IOException && cause.getCause() instanceof StorageException) {
+ return Pair.of(handleGetStorageException((StorageException) cause.getCause(), true), Option.empty());
+ }
+ throw e;
+ }
+ }
+
+ private @NotNull Pair<LockGetResult, Option<StorageLockFile>> getLockFileFromBlob(Blob blob) {
+ try (InputStream inputStream = Channels.newInputStream(blob.reader())) {
+ return Pair.of(LockGetResult.SUCCESS,
+ Option.of(StorageLockFile.createFromStream(inputStream, String.valueOf(blob.getGeneration()))));
+ } catch (IOException e) {
+ // Our createFromStream method does not throw IOExceptions, it wraps in HoodieIOException, however Sonar requires handling this.
+ throw new UncheckedIOException("Failed reading blob: " + lockFilePath, e);
+ }
+ }
+
+ @Override
+ public Option<String> readObject(String filePath, boolean checkExistsFirst) {
+ try {
+ // Parse the file path to get bucket and object path
+ Pair<String, String> bucketAndPath = StorageLockClient.parseBucketAndPath(filePath);
+ String bucket = bucketAndPath.getLeft();
+ String objectPath = bucketAndPath.getRight();
+
+ BlobId blobId = BlobId.of(bucket, objectPath);
+
+ if (checkExistsFirst) {
+ // First check if the file exists (lightweight metadata check)
+ Blob blob = gcsClient.get(blobId);
+
+ if (blob == null || !blob.exists()) {
+ // File doesn't exist - this is the common case for optional configs
+ logger.debug("JSON config file not found: {}", filePath);
+ return Option.empty();
+ }
+
+ // File exists, read its content
+ byte[] content = blob.getContent();
+ return Option.of(new String(content, UTF_8));
+ } else {
+ // Direct read without existence check
+ byte[] content = gcsClient.readAllBytes(blobId);
+ return Option.of(new String(content, UTF_8));
+ }
+ } catch (StorageException e) {
+ if (e.getCode() == NOT_FOUND_ERROR_CODE) {
+ logger.debug("JSON config file not found: {}", filePath);
+ } else {
+ logger.warn("Error reading JSON config file: {}", filePath, e);
+ }
+ return Option.empty();
+ } catch (Exception e) {
+ logger.warn("Error reading JSON config file: {}", filePath, e);
+ return Option.empty();
+ }
+ }
+
+ @Override
+ public boolean writeObject(String filePath, String content) {
+ try {
+ // Parse the file path to get bucket and object path
+ Pair<String, String> bucketAndPath = StorageLockClient.parseBucketAndPath(filePath);
+ String bucket = bucketAndPath.getLeft();
+ String objectPath = bucketAndPath.getRight();
+
+ BlobId blobId = BlobId.of(bucket, objectPath);
+ BlobInfo blobInfo = BlobInfo.newBuilder(blobId).build();
+
+ // Write the content to GCS
+ gcsClient.create(blobInfo, content.getBytes(UTF_8));
+
+ logger.debug("Successfully wrote object to: {}", filePath);
+ return true;
+ } catch (Exception e) {
+ logger.warn("Error writing object to: {}", filePath, e);
+ return false;
+ }
+ }
+
+ @Override
+ public void close() throws Exception {
+ this.gcsClient.close();
+ }
+
+ private long getGenerationNumber(Option<StorageLockFile> file) {
+ return (file.isPresent())
+ ? Long.parseLong(file.get().getVersionId())
+ : 0;
+ }
+}
\ No newline at end of file
diff --git a/hudi-gcp/src/test/java/org/apache/hudi/gcp/transaction/lock/TestGCSStorageLockClient.java b/hudi-gcp/src/test/java/org/apache/hudi/gcp/transaction/lock/TestGCSStorageLockClient.java
new file mode 100644
index 000000000000..8b270023d6c7
--- /dev/null
+++ b/hudi-gcp/src/test/java/org/apache/hudi/gcp/transaction/lock/TestGCSStorageLockClient.java
@@ -0,0 +1,304 @@
+/*
+ * 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.hudi.gcp.transaction.lock;
+
+import org.apache.hudi.client.transaction.lock.models.LockGetResult;
+import org.apache.hudi.client.transaction.lock.models.LockUpsertResult;
+import org.apache.hudi.client.transaction.lock.models.StorageLockData;
+import org.apache.hudi.client.transaction.lock.models.StorageLockFile;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.exception.HoodieIOException;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.cloud.ReadChannel;
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.BlobId;
+import com.google.cloud.storage.BlobInfo;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.mockito.stubbing.Answer;
+import org.slf4j.Logger;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Properties;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.contains;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+/**
+ * Test class for GCSStorageLockClient.
+ */
+@ExtendWith(MockitoExtension.class)
+public class TestGCSStorageLockClient {
+
+ private static final String OWNER_ID = "ownerId";
+ private static final String LOCK_FILE_URI = "gs://bucket/lockFilePath";
+ private static final String LOCK_FILE_PATH = "lockFilePath";
+ private static final String BUCKET_NAME = "bucket";
+
+ @Mock
+ private Storage mockStorage;
+
+ @Mock
+ private Blob mockBlob;
+
+ @Mock
+ private Logger mockLogger;
+
+ private GCSStorageLockClient lockService;
+
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+ /**
+ * Mocks the given {@link Blob} so that when {@code blob.reader()} is called,
+ * it returns a {@link ReadChannel} whose first read supplies JSON
+ * for the provided {@link StorageLockData}, followed by EOF on subsequent reads.
+ */
+ public static void mockBlobReaderWithLockData(Blob mockBlob, StorageLockData data) throws IOException {
+ String json = OBJECT_MAPPER.writeValueAsString(data);
+ byte[] jsonBytes = json.getBytes();
+
+ ReadChannel mockReadChannel = mock(ReadChannel.class);
+
+ Answer<Integer> fillBufferWithJson = invocation -> {
+ ByteBuffer buffer = invocation.getArgument(0);
+ buffer.put(jsonBytes);
+ return jsonBytes.length;
+ };
+ doAnswer(fillBufferWithJson)
+ .doAnswer(invocation -> -1)
+ .when(mockReadChannel).read(any(ByteBuffer.class));
+
+ when(mockBlob.reader()).thenReturn(mockReadChannel);
+ }
+
+ @BeforeEach
+ void setUp() {
+ lockService = new GCSStorageLockClient(OWNER_ID, LOCK_FILE_URI, new Properties(), (a) -> mockStorage, mockLogger);
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_noPreviousLock_success() {
+ StorageLockData lockData = new StorageLockData(false, 123L, "test-owner");
+
+ when(mockStorage.create(
+ any(BlobInfo.class),
+ any(byte[].class),
+ any(Storage.BlobTargetOption.class))
+ ).thenReturn(mockBlob);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result = lockService.tryUpsertLockFile(lockData, Option.empty());
+
+ assertNotNull(result.getRight(), "Expected a valid StorageLockFile on success");
+ verify(mockStorage).create(
+ any(BlobInfo.class),
+ any(byte[].class),
+ eq(Storage.BlobTargetOption.generationMatch(0))
+ );
+ verifyNoMoreInteractions(mockLogger);
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_withPreviousLock_success() {
+ StorageLockData lockData = new StorageLockData(false, 999L, "existing-owner");
+ StorageLockFile previousLockFile = new StorageLockFile(lockData, "123"); // versionId=123
+
+ when(mockStorage.create(
+ any(BlobInfo.class),
+ any(byte[].class),
+ any(Storage.BlobTargetOption.class))
+ ).thenReturn(mockBlob);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result =
+ lockService.tryUpsertLockFile(lockData, Option.of(previousLockFile));
+
+ assertNotNull(result, "Expected a valid StorageLockFile on success");
+ verify(mockStorage).create(
+ any(BlobInfo.class),
+ any(byte[].class),
+ eq(Storage.BlobTargetOption.generationMatch(123L))
+ );
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_preconditionFailed() {
+ StorageLockData lockData = new StorageLockData(false, 999L, "owner");
+ StorageLockFile previousLockFile = new StorageLockFile(lockData, "123");
+
+ StorageException exception = new StorageException(412, "Precondition Failed");
+ when(mockStorage.create(any(BlobInfo.class), any(byte[].class), any(Storage.BlobTargetOption.class)))
+ .thenThrow(exception);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result = lockService.tryUpsertLockFile(lockData, Option.of(previousLockFile));
+
+ assertEquals(LockUpsertResult.ACQUIRED_BY_OTHERS, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Should return empty when a 412 occurs");
+ verify(mockLogger).info(contains("Unable to write new lock file. Another process has modified this lockfile"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_rateLimitExceeded() {
+ StorageLockData lockData = new StorageLockData(false, 999L, "owner");
+ StorageException exception = new StorageException(429, "Rate Limit Exceeded");
+ when(mockStorage.create(any(BlobInfo.class), any(byte[].class), any(Storage.BlobTargetOption.class)))
+ .thenThrow(exception);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result = lockService.tryUpsertLockFile(lockData, Option.empty());
+
+ assertEquals(LockUpsertResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Should return empty when a 429 occurs");
+ verify(mockLogger).warn(contains("Rate limit exceeded"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_serverError() {
+ StorageLockData lockData = new StorageLockData(false, 999L, "owner");
+ StorageException exception = new StorageException(503, "Service Unavailable");
+ when(mockStorage.create(any(BlobInfo.class), any(byte[].class), any(Storage.BlobTargetOption.class)))
+ .thenThrow(exception);
+
+ Pair<LockUpsertResult, Option<StorageLockFile>> result = lockService.tryUpsertLockFile(lockData, Option.empty());
+
+ assertEquals(LockUpsertResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Should return empty when a 5xx error occurs");
+ verify(mockLogger).warn(contains("GCS returned internal server error code"), eq(OWNER_ID), eq(LOCK_FILE_PATH), eq(exception));
+ }
+
+ @Test
+ void testTryCreateOrUpdateLockFile_unexpectedError() {
+ StorageLockData lockData = new StorageLockData(false, 999L, "owner");
+ StorageException exception = new StorageException(400, "Bad Request");
+ when(mockStorage.create(any(BlobInfo.class), any(byte[].class), any(Storage.BlobTargetOption.class)))
+ .thenThrow(exception);
+
+ assertThrows(StorageException.class, () ->
+ lockService.tryUpsertLockFile(lockData, Option.empty()),
+ "Expected the method to rethrow the exception"
+ );
+ }
+
+ @Test
+ void testGetCurrentLockFile_blobNotFound() {
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenReturn(null);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+
+ assertEquals(LockGetResult.NOT_EXISTS, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Expected empty when no blob is found");
+ }
+
+ @Test
+ void testGetCurrentLockFile_blobFound() throws IOException {
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenReturn(mockBlob);
+ mockBlobReaderWithLockData(mockBlob, new StorageLockData(false, 123L, "owner"));
+ when(mockBlob.getGeneration()).thenReturn(123L);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+
+ assertEquals(LockGetResult.SUCCESS, result.getLeft());
+ assertNotNull(result.getRight(), "Should return a StorageLockFile if blob is found");
+ assertEquals("123", result.getRight().get().getVersionId(), "Version ID should match blob generation");
+ }
+
+ @Test
+ void testGetCurrentLockFile_404Error() {
+ StorageException exception404 = new StorageException(404, "Not Found");
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenThrow(exception404);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+
+ assertEquals(LockGetResult.NOT_EXISTS, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Should return empty on 404 error");
+ verify(mockLogger).info(contains("Object not found"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testGetCurrentLockFile_rateLimit() {
+ StorageException exception429 = new StorageException(429, "Rate Limit Exceeded");
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenThrow(exception429);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+
+ assertEquals(LockGetResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Should return empty on 429 error");
+ verify(mockLogger).warn(contains("Rate limit exceeded"), eq(OWNER_ID), eq(LOCK_FILE_PATH));
+ }
+
+ @Test
+ void testGetCurrentLockFile_serverError() {
+ StorageException exception500 = new StorageException(503, "Service Unavailable");
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenThrow(exception500);
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+
+ assertEquals(LockGetResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Should return empty on 5xx error");
+ verify(mockLogger).warn(contains("GCS returned internal server error code"), eq(OWNER_ID), eq(LOCK_FILE_PATH), eq(exception500));
+ }
+
+ @Test
+ void testGetCurrentLockFile_unexpectedError() {
+ StorageException exception400 = new StorageException(400, "Bad Request");
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenThrow(exception400);
+
+ assertThrows(StorageException.class, () -> lockService.readCurrentLockFile(),
+ "Should rethrow unexpected errors");
+ }
+
+ @Test
+ void testGetCurrentLockFile_IOException() {
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenReturn(mockBlob);
+ when(mockBlob.reader()).thenThrow(new HoodieIOException("IO Error"));
+
+ assertThrows(HoodieIOException.class, () -> lockService.readCurrentLockFile());
+ }
+
+ @Test
+ void testGetCurrentLockFile_IOExceptionWrapping404() {
+ when(mockStorage.get(BlobId.of(BUCKET_NAME, LOCK_FILE_PATH))).thenReturn(mockBlob);
+ when(mockBlob.reader()).thenThrow(new HoodieIOException("IO Error", new IOException(new StorageException(404, "storage 404"))));
+
+ Pair<LockGetResult, Option<StorageLockFile>> result = lockService.readCurrentLockFile();
+ assertEquals(LockGetResult.UNKNOWN_ERROR, result.getLeft());
+ assertTrue(result.getRight().isEmpty(), "Should return empty on IO exception wrapping 404 error");
+ }
+
+ @Test
+ void testClose() throws Exception {
+ lockService.close();
+ verify(mockStorage).close();
+ }
+}
diff --git a/hudi-gcp/src/test/java/org/apache/hudi/gcp/transaction/lock/TestGCSStorageLockClientAuditOperations.java b/hudi-gcp/src/test/java/org/apache/hudi/gcp/transaction/lock/TestGCSStorageLockClientAuditOperations.java
new file mode 100644
index 000000000000..02ea0719d764
--- /dev/null
+++ b/hudi-gcp/src/test/java/org/apache/hudi/gcp/transaction/lock/TestGCSStorageLockClientAuditOperations.java
@@ -0,0 +1,273 @@
+/*
+ * 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.hudi.gcp.transaction.lock;
+
+import org.apache.hudi.common.util.Option;
+
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.BlobId;
+import com.google.cloud.storage.BlobInfo;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Properties;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.contains;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for GCSStorageLockClient audit operations (readObject and writeObject methods)
+ */
+public class TestGCSStorageLockClientAuditOperations {
+
+ private Storage mockGcsClient;
+ private Blob mockBlob;
+ private Logger mockLogger;
+ private GCSStorageLockClient lockClient;
+
+ @BeforeEach
+ void setUp() {
+ mockGcsClient = mock(Storage.class);
+ mockLogger = mock(Logger.class);
+ mockBlob = mock(Blob.class);
+ String lockFileUri = "gs://test-bucket/table/.hoodie/.locks/table_lock.json";
+ String ownerId = "test-owner";
+ lockClient = new GCSStorageLockClient(
+ ownerId,
+ lockFileUri,
+ new Properties(),
+ props -> mockGcsClient,
+ mockLogger);
+ }
+
+ @Test
+ void testReadConfigWithCheckExistsFirstFileNotFound() {
+ String configPath = "gs://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // get() returns null for non-existent blob
+ when(mockGcsClient.get(any(BlobId.class)))
+ .thenReturn(null);
+
+ Option<String> result = lockClient.readObject(configPath, true);
+
+ assertTrue(result.isEmpty());
+ // Should only call get() for existence check, not readAllBytes
+ verify(mockGcsClient, times(1)).get(any(BlobId.class));
+ verify(mockGcsClient, never()).readAllBytes(any(BlobId.class));
+ }
+
+ @Test
+ void testReadConfigWithCheckExistsFirstBlobExistsFalse() {
+ String configPath = "gs://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // get() returns blob but exists() is false
+ when(mockGcsClient.get(any(BlobId.class)))
+ .thenReturn(mockBlob);
+ when(mockBlob.exists()).thenReturn(false);
+
+ Option<String> result = lockClient.readObject(configPath, true);
+
+ assertTrue(result.isEmpty());
+ // Should only check existence, not read content
+ verify(mockGcsClient, times(1)).get(any(BlobId.class));
+ verify(mockBlob, times(1)).exists();
+ verify(mockBlob, never()).getContent();
+ }
+
+ @Test
+ void testReadConfigWithCheckExistsFirstFileExists() {
+ String configPath = "gs://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+ String expectedContent = "{\"STORAGE_LP_AUDIT_SERVICE_ENABLED\": true}";
+
+ // get() returns existing blob
+ when(mockGcsClient.get(any(BlobId.class)))
+ .thenReturn(mockBlob);
+ when(mockBlob.exists()).thenReturn(true);
+ when(mockBlob.getContent())
+ .thenReturn(expectedContent.getBytes(StandardCharsets.UTF_8));
+
+ Option<String> result = lockClient.readObject(configPath, true);
+
+ assertTrue(result.isPresent());
+ assertEquals(expectedContent, result.get());
+ // Should call get() for existence check and getContent() for reading
+ verify(mockGcsClient, times(1)).get(any(BlobId.class));
+ verify(mockBlob, times(1)).exists();
+ verify(mockBlob, times(1)).getContent();
+ }
+
+ @Test
+ void testReadConfigWithoutCheckExistsFirstFileNotFound() {
+ String configPath = "gs://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // Direct readAllBytes throws 404 exception
+ StorageException notFoundException = new StorageException(404, "Not Found");
+ when(mockGcsClient.readAllBytes(any(BlobId.class)))
+ .thenThrow(notFoundException);
+
+ Option<String> result = lockClient.readObject(configPath, false);
+
+ assertTrue(result.isEmpty());
+ // Should not call get(), only readAllBytes
+ verify(mockGcsClient, never()).get(any(BlobId.class));
+ verify(mockGcsClient, times(1)).readAllBytes(any(BlobId.class));
+ }
+
+ @Test
+ void testReadConfigWithoutCheckExistsFirstFileExists() {
+ String configPath = "gs://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+ String expectedContent = "{\"STORAGE_LP_AUDIT_SERVICE_ENABLED\": false}";
+
+ // Direct readAllBytes returns content
+ when(mockGcsClient.readAllBytes(any(BlobId.class)))
+ .thenReturn(expectedContent.getBytes(StandardCharsets.UTF_8));
+
+ Option<String> result = lockClient.readObject(configPath, false);
+
+ assertTrue(result.isPresent());
+ assertEquals(expectedContent, result.get());
+ // Should not call get(), only readAllBytes
+ verify(mockGcsClient, never()).get(any(BlobId.class));
+ verify(mockGcsClient, times(1)).readAllBytes(any(BlobId.class));
+ }
+
+ @Test
+ void testReadConfigWithCheckExistsFirstOtherGcsError() {
+ String configPath = "gs://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // get() throws non-404 error
+ StorageException serverError = new StorageException(500, "Internal Server Error");
+ when(mockGcsClient.get(any(BlobId.class)))
+ .thenThrow(serverError);
+
+ Option<String> result = lockClient.readObject(configPath, true);
+
+ assertTrue(result.isEmpty());
+ verify(mockGcsClient, times(1)).get(any(BlobId.class));
+ verify(mockGcsClient, never()).readAllBytes(any(BlobId.class));
+ }
+
+ @Test
+ void testReadConfigWithInvalidUri() {
+ String invalidPath = "not-a-valid-uri";
+
+ Option<String> result = lockClient.readObject(invalidPath, false);
+
+ assertTrue(result.isEmpty());
+ // Should not make any GCS calls due to URI parsing error
+ verify(mockGcsClient, never()).get(any(BlobId.class));
+ verify(mockGcsClient, never()).readAllBytes(any(BlobId.class));
+ }
+
+ @Test
+ void testReadConfigWithRateLimitError() {
+ String configPath = "gs://test-bucket/table/.hoodie/.locks/audit_enabled.json";
+
+ // readAllBytes returns rate limit error
+ StorageException rateLimitException = new StorageException(429, "Too Many Requests");
+ when(mockGcsClient.readAllBytes(any(BlobId.class)))
+ .thenThrow(rateLimitException);
+
+ Option<String> result = lockClient.readObject(configPath, false);
+
+ assertTrue(result.isEmpty());
+ verify(mockGcsClient, times(1)).readAllBytes(any(BlobId.class));
+ }
+
+ // ================================
+ // writeObject() tests
+ // ================================
+
+ @Test
+ void testWriteObject_success() {
+ String filePath = "gs://test-bucket/audit/test-audit.jsonl";
+ String content = "{\"test\": \"data\"}\n";
+ when(mockGcsClient.create(any(BlobInfo.class), any(byte[].class))).thenReturn(mockBlob);
+
+ boolean result = lockClient.writeObject(filePath, content);
+
+ assertTrue(result);
+ verify(mockGcsClient).create(any(BlobInfo.class), eq(content.getBytes(UTF_8)));
+ verify(mockLogger).debug("Successfully wrote object to: {}", filePath);
+ }
+
+ @Test
+ void testWriteObject_storageException() {
+ String filePath = "gs://test-bucket/audit/test-audit.jsonl";
+ String content = "{\"test\": \"data\"}\n";
+ StorageException storageException = new StorageException(500, "Internal Server Error");
+ when(mockGcsClient.create(any(BlobInfo.class), any(byte[].class))).thenThrow(storageException);
+
+ boolean result = lockClient.writeObject(filePath, content);
+
+ assertFalse(result);
+ verify(mockLogger).warn(contains("Error writing object to"), eq(filePath), eq(storageException));
+ }
+
+ @Test
+ void testWriteObject_invalidPath() {
+ String invalidPath = "invalid-path";
+ String content = "{\"test\": \"data\"}\n";
+
+ boolean result = lockClient.writeObject(invalidPath, content);
+
+ assertFalse(result);
+ verify(mockLogger).warn(contains("Error writing object to"), eq(invalidPath), any(Exception.class));
+ }
+
+ @Test
+ void testWriteObject_emptyContent() {
+ String filePath = "gs://test-bucket/audit/empty-content.jsonl";
+ String content = "";
+ when(mockGcsClient.create(any(BlobInfo.class), any(byte[].class))).thenReturn(mockBlob);
+
+ boolean result = lockClient.writeObject(filePath, content);
+
+ assertTrue(result);
+ verify(mockGcsClient).create(any(BlobInfo.class), eq(content.getBytes(UTF_8)));
+ verify(mockLogger).debug("Successfully wrote object to: {}", filePath);
+ }
+
+ @Test
+ void testWriteObject_rateLimitExceeded() {
+ String filePath = "gs://test-bucket/audit/test-audit.jsonl";
+ String content = "{\"test\": \"data\"}\n";
+ StorageException rateLimitException = new StorageException(429, "Rate Limit Exceeded");
+ when(mockGcsClient.create(any(BlobInfo.class), any(byte[].class))).thenThrow(rateLimitException);
+
+ boolean result = lockClient.writeObject(filePath, content);
+
+ assertFalse(result);
+ verify(mockLogger).warn(contains("Error writing object to"), eq(filePath), eq(rateLimitException));
+ }
+}
diff --git a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/BaseProcedure.scala b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/BaseProcedure.scala
index b06aea2ac58c..7f7b2a324cc8 100644
--- a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/BaseProcedure.scala
+++ b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/BaseProcedure.scala
@@ -18,9 +18,11 @@
package org.apache.spark.sql.hudi.command.procedures
import org.apache.hudi.HoodieCLIUtils
+import org.apache.hudi.common.table.HoodieTableMetaClient
import org.apache.hudi.config.{HoodieIndexConfig, HoodieWriteConfig}
import org.apache.hudi.exception.HoodieException
import org.apache.hudi.index.HoodieIndex.IndexType
+
import org.apache.spark.api.java.JavaSparkContext
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.InternalRow
@@ -95,4 +97,9 @@ abstract class BaseProcedure extends Procedure {
)
}
+ protected def createMetaClient(jsc: JavaSparkContext, basePath: String): HoodieTableMetaClient = {
+ HoodieTableMetaClient.builder
+ .setConf(jsc.hadoopConfiguration())
+ .setBasePath(basePath).build
+ }
}
diff --git a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/CleanupAuditLockProcedure.scala b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/CleanupAuditLockProcedure.scala
new file mode 100644
index 000000000000..b32c7f6ea200
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/CleanupAuditLockProcedure.scala
@@ -0,0 +1,190 @@
+/*
+ * 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.spark.sql.hudi.command.procedures
+
+import org.apache.hudi.client.transaction.lock.audit.StorageLockProviderAuditService
+import org.apache.hudi.util.JFunction
+
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.types.{DataTypes, Metadata, StructField, StructType}
+
+import java.util.function.Supplier
+
+/**
+ * Spark SQL procedure for cleaning up old audit lock files for Hudi tables.
+ *
+ * This procedure removes audit lock files that are older than a specified number
+ * of days. It supports dry-run mode to preview which files would be deleted
+ * without actually removing them.
+ *
+ * Usage:
+ * {{{
+ * CALL cleanup_audit_lock(table => 'my_table')
+ * CALL cleanup_audit_lock(path => '/path/to/table', dry_run => true, age_days => 30)
+ * }}}
+ *
+ * The procedure operates on audit files in:
+ * `{table_path}/.hoodie/.locks/audit/`
+ *
+ * @author Apache Hudi
+ * @since 1.0.0
+ */
+class CleanupAuditLockProcedure extends BaseProcedure with ProcedureBuilder {
+ private val PARAMETERS = Array[ProcedureParameter](
+ ProcedureParameter.optional(0, "table", DataTypes.StringType),
+ ProcedureParameter.optional(1, "path", DataTypes.StringType),
+ ProcedureParameter.optional(2, "dry_run", DataTypes.BooleanType, false),
+ ProcedureParameter.optional(3, "age_days", DataTypes.IntegerType, 7)
+ )
+
+ private val OUTPUT_TYPE = new StructType(Array[StructField](
+ StructField("table", DataTypes.StringType, nullable = false, Metadata.empty),
+ StructField("files_cleaned", DataTypes.IntegerType, nullable = false, Metadata.empty),
+ StructField("dry_run", DataTypes.BooleanType, nullable = false, Metadata.empty),
+ StructField("age_days", DataTypes.IntegerType, nullable = false, Metadata.empty),
+ StructField("message", DataTypes.StringType, nullable = false, Metadata.empty)
+ ))
+
+ /**
+ * Returns the procedure parameters definition.
+ *
+ * @return Array of parameters: table (optional String), path (optional String),
+ * dry_run (optional Boolean, default false), and age_days (optional Integer, default 7)
+ */
+ def parameters: Array[ProcedureParameter] = PARAMETERS
+
+ /**
+ * Returns the output schema for the procedure result.
+ *
+ * @return StructType containing table, files_cleaned, dry_run, age_days, and message columns
+ */
+ def outputType: StructType = OUTPUT_TYPE
+
+ /**
+ * Executes the audit lock cleanup procedure.
+ *
+ * @param args Procedure arguments containing table name or path, dry_run flag, and age threshold
+ * @return Sequence containing a single Row with cleanup results
+ * @throws IllegalArgumentException if neither table nor path is provided, or both are provided
+ */
+ override def call(args: ProcedureArgs): Seq[Row] = {
+ super.checkArgs(PARAMETERS, args)
+
+ val tableName = getArgValueOrDefault(args, PARAMETERS(0))
+ val tablePath = getArgValueOrDefault(args, PARAMETERS(1))
+ val dryRun = getArgValueOrDefault(args, PARAMETERS(2)).getOrElse(false).asInstanceOf[Boolean]
+ val ageDays = getArgValueOrDefault(args, PARAMETERS(3)).getOrElse(7).asInstanceOf[Int]
+
+ // Validate age_days is positive
+ if (ageDays <= 0) {
+ throw new IllegalArgumentException("age_days must be a positive integer")
+ }
+
+ // Get the base path using BaseProcedure helper (handles table/path validation)
+ val basePath: String = getBasePath(tableName, tablePath)
+ val metaClient = createMetaClient(jsc, basePath)
+
+ // Use table name if provided, otherwise extract from path
+ val displayName = tableName.map(_.asInstanceOf[String]).getOrElse(tablePath.get.asInstanceOf[String])
+
+ try {
+ val auditFolderPath = new Path(StorageLockProviderAuditService.getAuditFolderPath(basePath))
+ val fs = metaClient.getFs
+
+ // Check if audit folder exists
+ if (!fs.exists(auditFolderPath)) {
+ val message = "No audit folder found - nothing to cleanup"
+ Seq(Row(displayName, 0, dryRun, ageDays, message))
+ } else {
+
+ // Calculate cutoff timestamp (ageDays ago)
+ val cutoffTime = System.currentTimeMillis() - (ageDays * 24 * 60 * 60 * 1000L)
+
+ // List all files in audit folder and filter by modification time
+ val allFiles = fs.listStatus(auditFolderPath)
+ val auditFiles = allFiles.filter(pathInfo => pathInfo.isFile && pathInfo.getPath.getName.endsWith(".jsonl"))
+ val oldFiles = auditFiles.filter(_.getModificationTime < cutoffTime)
+
+ if (oldFiles.isEmpty) {
+ val message = s"No audit files older than $ageDays days found"
+ Seq(Row(displayName, 0, dryRun, ageDays, message))
+ } else {
+
+ val fileCount = oldFiles.size
+
+ if (dryRun) {
+ val message = s"Dry run: Would delete $fileCount audit files older than $ageDays days"
+ Seq(Row(displayName, fileCount, dryRun, ageDays, message))
+ } else {
+ // Actually delete the files
+ var deletedCount = 0
+ var failedCount = 0
+
+ oldFiles.foreach { pathInfo =>
+ try {
+ fs.delete(pathInfo.getPath, false)
+ deletedCount += 1
+ } catch {
+ case e: Exception =>
+ failedCount += 1
+ // Log the error but continue with other files
+ }
+ }
+
+ val message = if (failedCount == 0) {
+ s"Successfully deleted $deletedCount audit files older than $ageDays days"
+ } else {
+ s"Deleted $deletedCount audit files, failed to delete $failedCount files"
+ }
+
+ Seq(Row(displayName, deletedCount, dryRun, ageDays, message))
+ }
+ }
+ }
+ } catch {
+ case e: Exception =>
+ val errorMessage = s"Cleanup failed: ${e.getMessage}"
+ Seq(Row(displayName, 0, dryRun, ageDays, errorMessage))
+ }
+ }
+
+ /**
+ * Builds a new instance of the CleanupAuditLockProcedure.
+ *
+ * @return New CleanupAuditLockProcedure instance
+ */
+ override def build: Procedure = new CleanupAuditLockProcedure()
+}
+
+
+/**
+ * Companion object for CleanupAuditLockProcedure containing constants and factory methods.
+ */
+object CleanupAuditLockProcedure {
+ /** The name used to register and invoke this procedure */
+ val NAME = "cleanup_audit_lock"
+
+ /**
+ * Factory method to create procedure builder instances.
+ *
+ * @return Supplier that creates new CleanupAuditLockProcedure instances
+ */
+ def builder: Supplier[ProcedureBuilder] = JFunction.toJavaSupplier(() => new CleanupAuditLockProcedure())
+}
diff --git a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala
index ad63ddbb29ee..1e746d7249f3 100644
--- a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala
+++ b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala
@@ -89,6 +89,10 @@ object HoodieProcedures {
,(ShowTablePropertiesProcedure.NAME, ShowTablePropertiesProcedure.builder)
,(HelpProcedure.NAME, HelpProcedure.builder)
,(ArchiveCommitsProcedure.NAME, ArchiveCommitsProcedure.builder)
+ ,(SetAuditLockProcedure.NAME, SetAuditLockProcedure.builder)
+ ,(ShowAuditLockStatusProcedure.NAME, ShowAuditLockStatusProcedure.builder)
+ ,(ValidateAuditLockProcedure.NAME, ValidateAuditLockProcedure.builder)
+ ,(CleanupAuditLockProcedure.NAME, CleanupAuditLockProcedure.builder)
)
}
}
diff --git a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/SetAuditLockProcedure.scala b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/SetAuditLockProcedure.scala
new file mode 100644
index 000000000000..625e3928d18c
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/SetAuditLockProcedure.scala
@@ -0,0 +1,189 @@
+/*
+ * 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.spark.sql.hudi.command.procedures
+
+import org.apache.hudi.client.transaction.lock.StorageLockClient
+import org.apache.hudi.client.transaction.lock.audit.StorageLockProviderAuditService
+import org.apache.hudi.common.table.HoodieTableMetaClient
+import org.apache.hudi.util.JFunction
+
+import org.apache.hadoop.fs.Path
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.databind.node.ObjectNode
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.types.{DataTypes, Metadata, StructField, StructType}
+
+import java.util.function.Supplier
+
+import scala.util.{Failure, Success, Try}
+
+
+/**
+ * Spark SQL procedure for enabling or disabling lock audit logging for Hudi tables.
+ *
+ * This procedure allows users to control audit logging for storage lock operations through
+ * Spark SQL commands. When enabled, lock operations will generate audit logs in JSONL format
+ * that track lock lifecycle events.
+ *
+ * Usage:
+ * {{{
+ * CALL set_audit_lock(table => 'my_table', state => 'enabled')
+ * CALL set_audit_lock(path => '/path/to/table', state => 'disabled')
+ * }}}
+ *
+ * The procedure creates or updates an audit configuration file at:
+ * `{table_path}/.hoodie/.locks/audit_enabled.json`
+ *
+ * @author Apache Hudi
+ * @since 1.0.0
+ */
+class SetAuditLockProcedure extends BaseProcedure with ProcedureBuilder {
+ private val PARAMETERS = Array[ProcedureParameter](
+ ProcedureParameter.optional(0, "table", DataTypes.StringType),
+ ProcedureParameter.optional(1, "path", DataTypes.StringType),
+ ProcedureParameter.required(2, "state", DataTypes.StringType)
+ )
+
+ private val OUTPUT_TYPE = new StructType(Array[StructField](
+ StructField("table", DataTypes.StringType, nullable = false, Metadata.empty),
+ StructField("audit_state", DataTypes.StringType, nullable = false, Metadata.empty),
+ StructField("message", DataTypes.StringType, nullable = false, Metadata.empty)
+ ))
+
+ private val OBJECT_MAPPER = new ObjectMapper()
+
... 1829 lines suppressed ...