From 93f38311350a46df1ef77c413f48efe15bf96c50 Mon Sep 17 00:00:00 2001
From: Akshay Rai <akshayrai09@gmail.com>
Date: Mon, 3 Nov 2025 16:26:44 +0530
Subject: [PATCH] Update task status on ZK if acquire was successful to reset
 previous state

---
 .../java/com/linkedin/datastream/server/DatastreamTaskImpl.java  | 1 +
 1 file changed, 1 insertion(+)

diff --git a/datastream-server/src/main/java/com/linkedin/datastream/server/DatastreamTaskImpl.java b/datastream-server/src/main/java/com/linkedin/datastream/server/DatastreamTaskImpl.java
index 43a4929b6..040f76037 100644
--- a/datastream-server/src/main/java/com/linkedin/datastream/server/DatastreamTaskImpl.java
+++ b/datastream-server/src/main/java/com/linkedin/datastream/server/DatastreamTaskImpl.java
@@ -337,6 +337,7 @@ public void acquire(Duration timeout) {
       });
 
       _zkAdapter.acquireTask(this, timeout);
+      setStatus(DatastreamTaskStatus.ok("Task lock acquired"));
     } catch (Exception e) {
       LOG.error(String.format("Failed to acquire task: %s with dependencies: %s", this.getDatastreamTaskName(),
           String.join(",", this.getDependencies())), e);