HIVE-29651: Update ZookeeperExternalSessionsRegistryClient to handle …

[tanishq-chugh]

…multiple HiveServer2 instances submitting DAGs concurrently to available Tez External Sessions

What changes were proposed in this pull request?

This PR introduces a distributed locking mechanism to synchronize Tez session assignments across multiple HiveServer2 instances.

Why are the changes needed?

To prevent Execution errors, when multiple HS2 instances tend to submit DAGs concurrently to same tez AMs

Does this PR introduce any user-facing change?

No

How was this patch tested?

Manual Testing + added UTs

[tanishq-chugh]

Hi @ayushtkn , Could you please help with a review on this patch ?
Thanks!

[ayushtkn]

Thanx @tanishq-chugh have dropped some comments

[ayushtkn]

can you add a debug log here

[tanishq-chugh]

Addressed this in commit - 7f6b08f

[ayushtkn]

this is wrong, known anti-patern, this will screw up if endTimeNs goes -ve due to overflow

[tanishq-chugh]

Right, changed the code to fix this in commit : fdead45

[ayushtkn]

Too Much Information :-), Please change it to debug

[tanishq-chugh]

Made the change in commit: 422c01f

[ayushtkn]

how do u reach to this magic number 15 here and 5 above? it is by default 60, why we need to change it?

[tanishq-chugh]

Leftover, removed in commit: 6bab8dc

[ayushtkn]

this will lead to flaky behaviour and neither guarantees FIFO. Ideally should use latches

something like

CountDownLatch r1Started = new CountDownLatch(1);
CountDownLatch r2Started = new CountDownLatch(1);

Future<String> future1 = executor.submit(() -> {
  r1Started.countDown();
  return registry1.getSession();
});

r1Started.await();

Future<String> future2 = executor.submit(() -> {
  r2Started.countDown();
  return registry2.getSession();
});

r2Started.await();

[tanishq-chugh]

Right, changed the test case in commit: 6bab8dc

[ayushtkn]

currious about the connection events, are we sure, if we loose connection and then CONNECTION_RECONNECTED is sent, are we sure the cache will replay all the missed events and our state would be correct?

Even more curios about the Connection Lost case
The network was down longer than the session timeout. Zookeeper deleted all of your ephemeral claim nodes. If you don't handle LOST, your local taken set will think it still owns the Tez AMs, but other HiveServer2 instances will see the nodes disappear and claim them right out from under you

[tanishq-chugh]

Nice catch @ayushtkn !
In case of CONNECTION_RECONNECTED , the cache does replay all the missed events, but while testing encountered a race condition between two listeners. Have addressed the same in commit : 26ef308

Regarding Connection Lost case, have added logic to kill running DAGs & clear taken state by the particular HS2 at the same time when ZK deletes its ephemeral claim nodes, in the same commit: 26ef308

[ayushtkn]

I am thinking do we need this? Was something like this possible

      CuratorCacheListener claimsListener = CuratorCacheListener.builder().forCreates(
          childData -> {
        if (childData == null)
          return;
        String applicationId = getApplicationId(childData);
        synchronized (lock) {
          available.remove(applicationId);
        }
      }).forDeletes(
          childData -> {
        if (childData == null)
          return;
        String applicationId = getApplicationId(childData);
        synchronized (lock) {
          if (!taken.contains(applicationId)) {
            available.add(applicationId);
            lock.notifyAll();
          }
        }
      }).build();

[tanishq-chugh]

Yes, changed this in commit: 484b2e4

[sonarqubecloud]

Quality Gate passed

Issues
2 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[Copilot]

Pull request overview

This PR aims to prevent multiple HiveServer2 instances from concurrently submitting DAGs to the same Tez external session by introducing ZooKeeper-based distributed coordination (a global queue + per-session claim znodes), and by routing DAG submission through a TezSession.submitDAG abstraction (with special handling for external sessions).

Changes:

• Add a distributed locking/claim mechanism in ZookeeperExternalSessionsRegistryClient to ensure each external Tez AM is claimed by at most one HS2 at a time.
• Add submitDAG(DAG) to TezSession and route Tez DAG submission through the session abstraction (enabling external-session-specific retry/cleanup behavior).
• Add/adjust unit tests to cover multi-client session claiming and FIFO behavior.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/ZookeeperExternalSessionsRegistryClient.java	Introduces global mutex queue + ephemeral claim znodes; adds connection-lost handling and claim cache.
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezSession.java	Extends session interface with submitDAG(DAG) to centralize submission behavior.
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezSessionState.java	Implements submitDAG by delegating to the underlying TezClient.
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezSessionPoolSession.java	Implements submitDAG by delegating through the wrapped base session.
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java	Switches DAG submission to sessionState.submitDAG(dag).
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezExternalSessionState.java	Adds external-session submit behavior (retry after killing running/orphan DAGs).
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezSessionPoolManager.java	Adds external-session handling on return-to-pool path (noted concurrency concern).
ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/TezJobMonitor.java	Adds helper to kill running DAGs for a specific application/session id.
ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestZookeeperExternalSessionsRegistryClient.java	Adds tests for multi-client claims and mutex FIFO queue behavior.
ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestTezTask.java	Updates mocks/verifications to match the new submitDAG call path.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.