[HUDI-6503] Make TableServiceClient's txnManager consistent with Writ…

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/BaseHoodieClient.java

@@ -76,26 +76,24 @@ public abstract class BaseHoodieClient implements Serializable, AutoCloseable {
    * incremental fashion.
    */
   private transient Option<EmbeddedTimelineService> timelineServer;
-  private final boolean shouldStopTimelineServer;
 
-  protected BaseHoodieClient(HoodieEngineContext context, HoodieWriteConfig clientConfig) {
-    this(context, clientConfig, Option.empty());
-  }
+  private boolean shouldCloseTxnManager;    // flag used for resource releasing
+  private boolean shouldStopTimelineServer; // flag used for resource releasing
 
-  protected BaseHoodieClient(HoodieEngineContext context, HoodieWriteConfig clientConfig,
-      Option<EmbeddedTimelineService> timelineServer) {
+  protected BaseHoodieClient(HoodieEngineContext context,
+                             HoodieWriteConfig clientConfig,
+                             Option<EmbeddedTimelineService> timelineServer,
+                             Option<TransactionManager> txnManager) {
     this.hadoopConf = context.getHadoopConf().get();
     this.fs = FSUtils.getFs(clientConfig.getBasePath(), hadoopConf);
     this.context = context;
     this.basePath = clientConfig.getBasePath();
     this.config = clientConfig;
-    this.timelineServer = timelineServer;
-    shouldStopTimelineServer = !timelineServer.isPresent();
     this.heartbeatClient = new HoodieHeartbeatClient(this.fs, this.basePath,
         clientConfig.getHoodieClientHeartbeatIntervalInMs(), clientConfig.getHoodieClientHeartbeatTolerableMisses());
     this.metrics = new HoodieMetrics(config);
-    this.txnManager = new TransactionManager(config, fs);
-    startEmbeddedServerView();
+    this.txnManager = initTxnManager(txnManager);
+    this.timelineServer = startEmbeddedServerView(timelineServer);
     initWrapperFSMetrics();
     runClientInitCallbacks();
   }
@@ -108,7 +106,9 @@ public void close() {
     stopEmbeddedServerView(true);
     this.context.setJobStatus("", "");
     this.heartbeatClient.close();
-    this.txnManager.close();
+    if (this.shouldCloseTxnManager) {
+      this.txnManager.close();
+    }
   }
 
   private synchronized void stopEmbeddedServerView(boolean resetViewStorageConfig) {
@@ -125,21 +125,31 @@ private synchronized void stopEmbeddedServerView(boolean resetViewStorageConfig)
     }
   }
 
-  private synchronized void startEmbeddedServerView() {
-    if (config.isEmbeddedTimelineServerEnabled()) {
-      if (!timelineServer.isPresent()) {
+  private synchronized Option<EmbeddedTimelineService> startEmbeddedServerView(Option<EmbeddedTimelineService> reused) {
+    if (reused.isPresent()) {
+      LOG.info("Timeline Server already running. Not restarting the service");
+      return reused;
+    } else {
+      if (config.isEmbeddedTimelineServerEnabled()) {
         // Run Embedded Timeline Server
         try {
-          timelineServer = EmbeddedTimelineServerHelper.createEmbeddedTimelineService(context, config);
+          this.shouldStopTimelineServer = true;
+          return EmbeddedTimelineServerHelper.createEmbeddedTimelineService(context, config);
         } catch (IOException e) {
           LOG.warn("Unable to start timeline service. Proceeding as if embedded server is disabled", e);
           stopEmbeddedServerView(false);
         }
-      } else {
-        LOG.info("Timeline Server already running. Not restarting the service");
       }
+    }
+    return Option.empty();
+  }
+
+  private TransactionManager initTxnManager(Option<TransactionManager> reused) {
+    if (reused.isPresent()) {
+      return reused.get();
     } else {
-      LOG.info("Embedded Timeline Server is disabled. Not starting timeline service");
+      this.shouldCloseTxnManager = true;
+      return new TransactionManager(config, fs);
     }
   }
 
@@ -186,6 +196,10 @@ public HoodieHeartbeatClient getHeartbeatClient() {
     return heartbeatClient;
   }
 
+  public Option<TransactionManager> getTxnManager() {
+    return Option.of(txnManager);
+  }
+
   /**
    * Resolve write conflicts before commit.
    *

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/BaseHoodieTableServiceClient.java

@@ -28,6 +28,7 @@
 import org.apache.hudi.avro.model.HoodieRollbackPlan;
 import org.apache.hudi.client.embedded.EmbeddedTimelineService;
 import org.apache.hudi.client.heartbeat.HeartbeatUtils;
+import org.apache.hudi.client.transaction.TransactionManager;
 import org.apache.hudi.common.HoodiePendingRollbackInfo;
 import org.apache.hudi.common.data.HoodieData;
 import org.apache.hudi.common.engine.HoodieEngineContext;
@@ -93,8 +94,9 @@ public abstract class BaseHoodieTableServiceClient<O> extends BaseHoodieClient i
 
   protected BaseHoodieTableServiceClient(HoodieEngineContext context,
                                          HoodieWriteConfig clientConfig,
-                                         Option<EmbeddedTimelineService> timelineService) {
-    super(context, clientConfig, timelineService);
+                                         Option<EmbeddedTimelineService> timelineService,
+                                         Option<TransactionManager> txnManager) {
+    super(context, clientConfig, timelineService, txnManager);
   }
 
   protected void startAsyncCleanerService(BaseHoodieWriteClient writeClient) {
@@ -172,7 +174,7 @@ private void inlineCompaction(HoodieTable table, Option<Map<String, String>> ext
   }
 
   /**
-   * Ensures compaction instant is in expected state and performs Log Compaction for the workload stored in instant-time.s
+   * Ensures compaction instant is in expected state and performs Log Compaction for the workload stored in instant-time.
    *
    * @param compactionInstantTime Compaction Instant Time
    * @return Collection of Write Status

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/BaseHoodieWriteClient.java

@@ -136,32 +136,18 @@ public abstract class BaseHoodieWriteClient<T, I, K, O> extends BaseHoodieClient
 
   protected BaseHoodieTableServiceClient<O> tableServiceClient;
 
-  /**
-   * Create a write client, with new hudi index.
-   * @param context HoodieEngineContext
-   * @param writeConfig instance of HoodieWriteConfig
-   * @param upgradeDowngradeHelper engine-specific instance of {@link SupportsUpgradeDowngrade}
-   */
-  @Deprecated
-  public BaseHoodieWriteClient(HoodieEngineContext context,
-                               HoodieWriteConfig writeConfig,
-                               SupportsUpgradeDowngrade upgradeDowngradeHelper) {
-    this(context, writeConfig, Option.empty(), upgradeDowngradeHelper);
-  }
-
   /**
    * Create a write client, allows to specify all parameters.
    *
    * @param context         HoodieEngineContext
    * @param writeConfig     instance of HoodieWriteConfig
    * @param timelineService Timeline Service that runs as part of write client.
    */
-  @Deprecated
   public BaseHoodieWriteClient(HoodieEngineContext context,
                                HoodieWriteConfig writeConfig,
                                Option<EmbeddedTimelineService> timelineService,
                                SupportsUpgradeDowngrade upgradeDowngradeHelper) {
-    super(context, writeConfig, timelineService);
+    super(context, writeConfig, timelineService, Option.empty());
     this.index = createIndex(writeConfig);
     this.upgradeDowngradeHelper = upgradeDowngradeHelper;
   }
@@ -1311,12 +1297,13 @@ protected void releaseResources(String instantTime) {
 
   @Override
   public void close() {
+    // Close table service client first
+    this.tableServiceClient.close();
     // Stop timeline-server if running
     super.close();
     // Calling this here releases any resources used by your index, so make sure to finish any related operations
     // before this point
     this.index.close();
-    this.tableServiceClient.close();
   }
 
   private void setWriteTimer(HoodieTable table) {

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/CompactionAdminClient.java

@@ -67,7 +67,7 @@ public class CompactionAdminClient extends BaseHoodieClient {
   private static final Logger LOG = LoggerFactory.getLogger(CompactionAdminClient.class);
 
   public CompactionAdminClient(HoodieEngineContext context, String basePath) {
-    super(context, HoodieWriteConfig.newBuilder().withPath(basePath).build());
+    super(context, HoodieWriteConfig.newBuilder().withPath(basePath).build(), Option.empty(), Option.empty());
   }
 
   /**

hudi-client/hudi-client-common/src/main/java/org/apache/hudi/client/transaction/TransactionManager.java

@@ -84,10 +84,8 @@ protected synchronized boolean reset(Option<HoodieInstant> callerInstant,
   }
 
   public void close() {
-    if (isLockRequired) {
-      lockManager.close();
-      LOG.info("Transaction manager closed");
-    }
+    lockManager.close();
+    LOG.info("Transaction manager closed");
   }
 
   public LockManager getLockManager() {

hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/client/HoodieFlinkTableServiceClient.java

@@ -20,6 +20,7 @@
 
 import org.apache.hudi.client.common.HoodieFlinkEngineContext;
 import org.apache.hudi.client.embedded.EmbeddedTimelineService;
+import org.apache.hudi.client.transaction.TransactionManager;
 import org.apache.hudi.common.data.HoodieData;
 import org.apache.hudi.common.engine.HoodieEngineContext;
 import org.apache.hudi.common.model.HoodieCommitMetadata;
@@ -68,8 +69,9 @@ public class HoodieFlinkTableServiceClient<T> extends BaseHoodieTableServiceClie
 
   protected HoodieFlinkTableServiceClient(HoodieEngineContext context,
                                           HoodieWriteConfig clientConfig,
-                                          Option<EmbeddedTimelineService> timelineService) {
-    super(context, clientConfig, timelineService);
+                                          Option<EmbeddedTimelineService> timelineService,
+                                          Option<TransactionManager> txnManager) {
+    super(context, clientConfig, timelineService, txnManager);
   }
 
   @Override

hudi-client/hudi-flink-client/src/main/java/org/apache/hudi/client/HoodieFlinkWriteClient.java

@@ -86,9 +86,9 @@ public class HoodieFlinkWriteClient<T> extends
   private final Map<String, Path> bucketToHandles;
 
   public HoodieFlinkWriteClient(HoodieEngineContext context, HoodieWriteConfig writeConfig) {
-    super(context, writeConfig, FlinkUpgradeDowngradeHelper.getInstance());
+    super(context, writeConfig, Option.empty(), FlinkUpgradeDowngradeHelper.getInstance());
     this.bucketToHandles = new HashMap<>();
-    this.tableServiceClient = new HoodieFlinkTableServiceClient<>(context, writeConfig, getTimelineServer());
+    this.tableServiceClient = new HoodieFlinkTableServiceClient<>(context, writeConfig, getTimelineServer(), getTxnManager());
   }
 
   /**

hudi-client/hudi-flink-client/src/test/java/org/apache/hudi/client/TestFlinkWriteClient.java

@@ -23,13 +23,15 @@
 import org.apache.hudi.testutils.HoodieFlinkClientTestHarness;
 
 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 static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 
 public class TestFlinkWriteClient extends HoodieFlinkClientTestHarness {
 
@@ -61,4 +63,22 @@ public void testWriteClientAndTableServiceClientWithTimelineServer(
 
     writeClient.close();
   }
+
+  @Test
+  public void testWriteClientAndTableServiceClientWithTxnManager() {
+    HoodieWriteConfig writeConfig = HoodieWriteConfig.newBuilder()
+        .withPath(metaClient.getBasePathV2().toString())
+        .build();
+
+    HoodieFlinkWriteClient writeClient = new HoodieFlinkWriteClient(context, writeConfig);
+
+    assertNotNull(writeClient.getTxnManager());
+
+    assertEquals(
+        writeClient.getTxnManager(),
+        writeClient.getTableServiceClient().getTxnManager()
+    );
+
+    writeClient.close();
+  }
 }

hudi-client/hudi-java-client/src/main/java/org/apache/hudi/client/HoodieJavaTableServiceClient.java

@@ -19,6 +19,7 @@
 package org.apache.hudi.client;
 
 import org.apache.hudi.client.embedded.EmbeddedTimelineService;
+import org.apache.hudi.client.transaction.TransactionManager;
 import org.apache.hudi.common.engine.HoodieEngineContext;
 import org.apache.hudi.common.model.HoodieCommitMetadata;
 import org.apache.hudi.common.util.Option;
@@ -37,8 +38,9 @@ public class HoodieJavaTableServiceClient extends BaseHoodieTableServiceClient<L
 
   protected HoodieJavaTableServiceClient(HoodieEngineContext context,
                                          HoodieWriteConfig clientConfig,
-                                         Option<EmbeddedTimelineService> timelineService) {
-    super(context, clientConfig, timelineService);
+                                         Option<EmbeddedTimelineService> timelineService,
+                                         Option<TransactionManager> txnManager) {
+    super(context, clientConfig, timelineService, txnManager);
   }
 
   @Override

hudi-client/hudi-java-client/src/main/java/org/apache/hudi/client/HoodieJavaWriteClient.java

@@ -18,7 +18,6 @@
 
 package org.apache.hudi.client;
 
-import org.apache.hudi.client.common.HoodieJavaEngineContext;
 import org.apache.hudi.client.embedded.EmbeddedTimelineService;
 import org.apache.hudi.common.data.HoodieListData;
 import org.apache.hudi.common.engine.HoodieEngineContext;
@@ -52,22 +51,20 @@ public class HoodieJavaWriteClient<T> extends
     BaseHoodieWriteClient<T, List<HoodieRecord<T>>, List<HoodieKey>, List<WriteStatus>> {
 
   public HoodieJavaWriteClient(HoodieEngineContext context, HoodieWriteConfig writeConfig) {
-    super(context, writeConfig, JavaUpgradeDowngradeHelper.getInstance());
-    this.tableServiceClient = new HoodieJavaTableServiceClient(context, writeConfig, getTimelineServer());
+    this(context, writeConfig, Option.empty());
   }
 
   public HoodieJavaWriteClient(HoodieEngineContext context,
                                HoodieWriteConfig writeConfig,
-                               boolean rollbackPending,
                                Option<EmbeddedTimelineService> timelineService) {
     super(context, writeConfig, timelineService, JavaUpgradeDowngradeHelper.getInstance());
-    this.tableServiceClient = new HoodieJavaTableServiceClient(context, writeConfig, getTimelineServer());
+    this.tableServiceClient = new HoodieJavaTableServiceClient(context, writeConfig, getTimelineServer(), getTxnManager());
   }
 
   @Override
   public List<HoodieRecord<T>> filterExists(List<HoodieRecord<T>> hoodieRecords) {
     // Create a Hoodie table which encapsulated the commits and files visible
-    HoodieJavaTable<T> table = HoodieJavaTable.create(config, (HoodieJavaEngineContext) context);
+    HoodieJavaTable<T> table = HoodieJavaTable.create(config, context);
     Timer.Context indexTimer = metrics.getIndexCtx();
     List<HoodieRecord<T>> recordsWithLocation = getIndex().tagLocation(HoodieListData.eager(hoodieRecords), context, table).collectAsList();
     metrics.updateIndexMetrics(LOOKUP_STR, metrics.getDurationInMs(indexTimer == null ? 0L : indexTimer.stop()));
@@ -239,7 +236,7 @@ protected void completeCompaction(HoodieCommitMetadata metadata,
 
   @Override
   protected HoodieWriteMetadata<List<WriteStatus>> compact(String compactionInstantTime,
-                                      boolean shouldComplete) {
+                                                           boolean shouldComplete) {
     throw new HoodieNotSupportedException("Compact is not supported in HoodieJavaClient");
   }
 

hudi-client/hudi-java-client/src/test/java/org/apache/hudi/client/TestHoodieJavaWriteClientInsert.java

@@ -61,6 +61,7 @@
 import static org.apache.hudi.index.HoodieIndex.IndexType.INMEMORY;
 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;
 
 public class TestHoodieJavaWriteClientInsert extends HoodieJavaClientTestHarness {
@@ -120,7 +121,7 @@ public void testWriteClientAndTableServiceClientWithTimelineServer(
           new EmbeddedTimelineService(context, null, writeConfig);
       timelineService.startServer();
       writeConfig.setViewStorageConfig(timelineService.getRemoteFileSystemViewConfig());
-      writeClient = new HoodieJavaWriteClient(context, writeConfig, true, Option.of(timelineService));
+      writeClient = new HoodieJavaWriteClient(context, writeConfig, Option.of(timelineService));
       // Both the write client and the table service client should use the same passed-in
       // timeline server instance.
       assertEquals(timelineService, writeClient.getTimelineServer().get());
@@ -142,6 +143,25 @@ public void testWriteClientAndTableServiceClientWithTimelineServer(
     writeClient.close();
   }
 
+  @Test
+  public void testWriteClientAndTableServiceClientWithTxnManager() {
+    HoodieWriteConfig writeConfig = HoodieWriteConfig.newBuilder()
+        .withPath(metaClient.getBasePathV2().toString())
+        .withIndexConfig(HoodieIndexConfig.newBuilder().withIndexType(INMEMORY).build())
+        .build();
+
+    HoodieJavaWriteClient writeClient = new HoodieJavaWriteClient(context, writeConfig);
+
+    assertNotNull(writeClient.getTxnManager());
+
+    assertEquals(
+        writeClient.getTxnManager(),
+        writeClient.getTableServiceClient().getTxnManager()
+    );
+
+    writeClient.close();
+  }
+
   @Test
   public void testInsert() throws Exception {
     HoodieWriteConfig config = makeHoodieClientConfigBuilder(basePath).withMergeAllowDuplicateOnInserts(true).build();

hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/client/SparkRDDTableServiceClient.java

@@ -21,6 +21,7 @@
 import org.apache.hudi.avro.model.HoodieClusteringGroup;
 import org.apache.hudi.avro.model.HoodieClusteringPlan;
 import org.apache.hudi.client.embedded.EmbeddedTimelineService;
+import org.apache.hudi.client.transaction.TransactionManager;
 import org.apache.hudi.common.data.HoodieData;
 import org.apache.hudi.common.engine.HoodieEngineContext;
 import org.apache.hudi.common.model.HoodieCommitMetadata;
@@ -67,8 +68,9 @@ public class SparkRDDTableServiceClient<T> extends BaseHoodieTableServiceClient<
 
   protected SparkRDDTableServiceClient(HoodieEngineContext context,
                                        HoodieWriteConfig clientConfig,
-                                       Option<EmbeddedTimelineService> timelineService) {
-    super(context, clientConfig, timelineService);
+                                       Option<EmbeddedTimelineService> timelineService,
+                                       Option<TransactionManager> txnManager) {
+    super(context, clientConfig, timelineService, txnManager);
   }
 
   @Override