feat(rbac): Netty Proxy & Sticky Lease Manager for Scalable Multi-Tenant Task Workers

Dockerfile

@@ -0,0 +1,28 @@
+# Copyright © 2026 Cask Data, Inc.
+#
+# Licensed 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.
+
+FROM us-east1-docker.pkg.dev/j145774183a931adb-tp/cdf-dev-shru/cloud-data-fusion:latest
+# For OSS CDAP, use "FROM gcr.io/cdapio/cdap:latest"
+
+RUN rm -rf /opt/cdap/master/lib/io.cdap.cdap.cdap-common-6.12.0-SNAPSHOT.jar
+
+COPY cdap-common/target/cdap-common-6.12.0-SNAPSHOT.jar /opt/cdap/master/lib/io.cdap.cdap.cdap-common-6.12.0-SNAPSHOT.jar
+
+RUN rm -rf /opt/cdap/master/ext/environments/k8s/io.cdap.cdap.cdap-kubernetes-6.12.0-SNAPSHOT.jar
+
+
+COPY cdap-kubernetes/target/cdap-kubernetes-6.12.0-SNAPSHOT.jar /opt/cdap/master/ext/environments/k8s/io.cdap.cdap.cdap-kubernetes-6.12.0-SNAPSHOT.jar
+
+
+RUN chmod -R 755 /opt/cdap

_agents/rules/rbac_taskmanager.md

@@ -0,0 +1,18 @@
+# CDAP RBAC Warm Sticky Leases & Task Manager Service
+
+When working on the RBAC Everywhere feature, Namespaced Service Accounts (NSA), or Task Worker pod scaling in this repository:
+
+1.  **Architecture Context**:
+    *   Refer to the detailed Warm Sticky Lease research notes here:
+        [rbac_taskmanager_research_notes.md](file:///usr/local/google/home/venkataramansh/.gemini/jetski/brain/2099d0c8-9e2b-4db5-a81a-5cfb437c1660/rbac_taskmanager_research_notes.md)
+    *   This feature resolves the "429 collision storm" and cold-start latencies (~40s) by shifting tenant isolation from the request level to the namespace/pod lease level.
+
+2.  **Core Routing Rules**:
+    *   **Direct Routing**: `RemoteClient` must bypass K8s round-robin load balancing. It resolves the headless task-worker service via DNS expansion to individual pod IPs, queries the `TaskManager` service to resolve the lease, and routes directly to the leased pod IP.
+    *   **Lease Registry**: The `TaskManager` service holds the lease maps in-memory to prevent Spanner database write contention. It uses a `ReentrantLock` to serialize checks/actions and prevent race conditions.
+    *   **Local Guard**: Individual Task Workers use `StickyLeaseManager` as a fail-safe to reject mismatching namespace requests locally with a `429`.
+
+3.  **Key Classes**:
+    *   App Fabric / Client: `RemoteClient`, `RemoteTaskExecutor`, `KubeDiscoveryService`
+    *   Task Manager: `TaskManager`, `TaskManagerHttpHandler`, `TaskManagerService`, `TaskManagerMain`
+    *   Task Worker: `StickyLeaseManager`, `TaskWorkerHttpHandlerInternal`

_agents/rules/task_manager_context.md

@@ -0,0 +1,19 @@
+# Task Manager & Sticky Lease Context
+
+## Architecture
+*   **Environment:** Cloud Data Fusion (CDF) on GKE (ZooKeeper-free).
+*   **Service Discovery:** Headless DNS expansion to pod IPs.
+*   **Coordination:** Standalone single-replica HTTP `TaskManager` service.
+*   **Local Guard:** `StickyLeaseManager` on Task Worker pods enforcing "First-Write Wins" lease lock.
+
+## Concurrency & Workload Limits
+*   **Concurrency:** Max 10 concurrent tasks per pod (enforced by `podActiveTaskCounts` / `activeTasks`).
+*   **Lifetime Limit:** Max 10 tasks before reset (enforced by `podTotalTaskProcessedCounts`).
+    *   *Increment in `resolve`:* Enforces strict limit of 10 tasks started (safe, current behavior).
+    *   *Increment in `finish`:* Allows better utilization but pod can process up to 19 tasks due to concurrency.
+
+## Reliability & Recovery
+*   **Downtime:** `RemoteClient` falls back to local consistent hashing if `TaskManager` is down.
+*   **State Recovery:** To recover from `TaskManager` restarts without polling, use a self-correction pattern:
+    *   Task Worker returns `409 Conflict` (with active namespace in body) on lease mismatch.
+    *   `RemoteClient` parses 409 and notifies `TaskManager` to update its lease map.

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/worker/sidecar/ArtifactLocalizerClient.java

@@ -103,7 +103,17 @@ private File sendRequest(ArtifactId artifactId, boolean unpack)
     Multimap<String, String> headers = httpRequest.getHeaders();
     internalAuthenticator.applyInternalAuthenticationHeaders(headers::put);
 
-    HttpResponse httpResponse = HttpRequests.execute(httpRequest);
+    long startTime = System.currentTimeMillis();
+    HttpResponse httpResponse;
+    try {
+      httpResponse = HttpRequests.execute(httpRequest);
+    } catch (IOException e) {
+      LOG.error("shruzard ArtifactLocalizerClient request for {} failed after {} ms with error: {}",
+          artifactId, System.currentTimeMillis() - startTime, e.getMessage());
+      throw e;
+    }
+    LOG.info("shruzard ArtifactLocalizerClient request for {} completed in {} ms (response code: {})",
+        artifactId, System.currentTimeMillis() - startTime, httpResponse.getResponseCode());
 
     if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) {
       if (httpResponse.getResponseCode() == HttpURLConnection.HTTP_NOT_FOUND) {

cdap-common/src/main/java/io/cdap/cdap/common/conf/Constants.java

@@ -142,6 +142,7 @@ public static final class Service {
     public static final String INTERNAL_ROUTER = "internal.router";
     public static final String AUTHENTICATION = "authentication";
     public static final String TASK_WORKER = "task.worker";
+    public static final String TASK_MANAGER = "task.manager";
     public static final String SYSTEM_WORKER = "system.worker";
     public static final String ARTIFACT_LOCALIZER = "artifact.localizer";
     public static final String SYSTEM_METRICS_EXPORTER = "system.metrics.exporter";

cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright © 2026 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.cdap.common.internal.remote;
+
+/**
+ * Tracks the routing state and load for a given worker pod IP.
+ */
+public class PodState {
+    private String leasedNamespace;
+    private int inflightRequests;
+    private long lastActivityTime;
+
+    public PodState(String leasedNamespace, int inflightRequests) {
+        this.leasedNamespace = leasedNamespace;
+        this.inflightRequests = inflightRequests;
+        // Subtract 40 seconds worth of nanos to instantly trigger predictions on boot
+        this.lastActivityTime = System.nanoTime() 
+            - java.util.concurrent.TimeUnit.SECONDS.toNanos(40);
+    }
+
+    public String getLeasedNamespace() {
+        return leasedNamespace;
+    }
+
+    public void setLeasedNamespace(String leasedNamespace) {
+        this.leasedNamespace = leasedNamespace;
+    }
+
+    public int getInflightRequests() {
+        return inflightRequests;
+    }
+
+    public void setInflightRequests(int inflightRequests) {
+        this.inflightRequests = inflightRequests;
+    }
+
+    public long getLastActivityTime() {
+        return lastActivityTime;
+    }
+
+    public void setLastActivityTime(long lastActivityTime) {
+        this.lastActivityTime = lastActivityTime;
+    }
+}

cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java

@@ -0,0 +1,126 @@
+/*
+ * Copyright © 2026 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.cdap.common.internal.remote;
+
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelFutureListener;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.handler.codec.http.HttpResponse;
+
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ProxyBackendHandler extends ChannelInboundHandlerAdapter {
+
+    private static final Logger LOG = LoggerFactory.getLogger(ProxyBackendHandler.class);
+
+
+    private final Channel inboundChannel;
+    private final Map<String, PodState> podRegistry;
+    private final String targetWorkerAddress;
+    private boolean decremented = false;
+
+    public ProxyBackendHandler(Channel inboundChannel, Map<String, PodState> podRegistry, String targetWorkerAddress) {
+        this.inboundChannel = inboundChannel;
+        this.podRegistry = podRegistry;
+        this.targetWorkerAddress = targetWorkerAddress;
+    }
+
+    private synchronized void decrementInflight() {
+        if (!decremented) {
+            PodState state = podRegistry.get(targetWorkerAddress);
+            if (state != null) {
+                synchronized (state) {
+                    state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1));
+                    state.setLastActivityTime(System.nanoTime());
+                }
+            }
+            decremented = true;
+        }
+    }
+
+    @Override
+    public void channelRead(ChannelHandlerContext ctx, Object msg) {
+        if (msg instanceof HttpResponse) {
+            HttpResponse resp = (HttpResponse) msg;
+            PodState state = podRegistry.get(targetWorkerAddress);
+            if (state != null) {
+                // Thread-safe update from Worker Ground Truth headers
+                synchronized (state) {
+                    String activeTasksStr = resp.headers().get("X-Active-Tasks");
+                    String leasedNamespace = resp.headers().get("X-Leased-Namespace");
+                    // Treat Task Worker as strict source of truth for load ONLY if it rejects us
+                    if (resp.status().code() == 429 || resp.status().code() == 409) {
+                        if (activeTasksStr != null) {
+                            state.setInflightRequests(Integer.parseInt(activeTasksStr));
+                        }
+                    }
+
+                    if (leasedNamespace != null) {
+                        state.setLeasedNamespace(leasedNamespace);
+                    }
+
+                    state.setLastActivityTime(System.nanoTime());
+
+                    if (activeTasksStr != null || leasedNamespace != null) {
+                        LOG.info("shruzard - ProxyBackendHandler: Header Sync "
+                                 + "PodState for {}. Local Occupancy: {}, Remote Tasks: {}, Namespace: {}", 
+                            targetWorkerAddress, state.getInflightRequests(), activeTasksStr, 
+                            state.getLeasedNamespace());
+                    }
+                }
+            }
+        }
+
+        if (msg instanceof io.netty.handler.codec.http.LastHttpContent) {
+            decrementInflight();
+        }
+
+        // Forward worker responses directly back to the client
+        inboundChannel.writeAndFlush(msg).addListener((ChannelFutureListener) future -> {
+            if (future.isSuccess()) {
+                ctx.channel().read();
+            } else {
+                future.channel().close();
+            }
+        });
+    }
+
+    @Override
+    public void channelWritabilityChanged(ChannelHandlerContext ctx) {
+        // Backend Worker channel is saturated; pause reading from App Fabric client
+        if (inboundChannel != null && inboundChannel.isActive()) {
+            inboundChannel.config().setAutoRead(ctx.channel().isWritable());
+        }
+        ctx.fireChannelWritabilityChanged();
+    }
+
+    @Override
+    public void channelInactive(ChannelHandlerContext ctx) {
+        decrementInflight();
+        ProxyFrontendHandler.closeOnFlush(inboundChannel);
+    }
+
+    @Override
+    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
+        decrementInflight();
+        cause.printStackTrace();
+        ProxyFrontendHandler.closeOnFlush(ctx.channel());
+    }
+}