From c9f58baecdc04d2b55c6f958f82c8e9f1d0079f4 Mon Sep 17 00:00:00 2001 From: sidhdirenge Date: Mon, 29 Jun 2026 09:09:46 +0000 Subject: [PATCH 01/54] feat(rbac-poc): implement standalone Task Manager Service and integrate into RemoteClient --- .../common/internal/remote/RemoteClient.java | 182 ++++++++++++++++- .../common/internal/remote/TaskManager.java | 185 ++++++++++++++++++ .../remote/TaskManagerHttpHandler.java | 152 ++++++++++++++ .../internal/remote/TaskManagerMain.java | 55 ++++++ .../internal/remote/TaskManagerTest.java | 90 +++++++++ task-manager-service.yaml | 49 +++++ 6 files changed, 704 insertions(+), 9 deletions(-) create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java create mode 100644 cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java create mode 100644 task-manager-service.yaml diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index af8a37222fed..7fb8bc8ee829 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -20,6 +20,7 @@ import com.google.common.collect.HashMultimap; import com.google.common.collect.Multimap; import com.google.common.net.HttpHeaders; +import com.google.gson.Gson; import io.cdap.cdap.api.retry.Idempotency; import io.cdap.cdap.api.retry.RetryableException; import io.cdap.cdap.api.service.ServiceUnavailableException; @@ -45,13 +46,18 @@ import java.net.MalformedURLException; import java.net.URI; import java.net.URL; +import java.util.ArrayList; +import java.util.Comparator; import java.util.EnumSet; +import java.util.List; import java.util.concurrent.TimeUnit; import java.util.function.BiConsumer; import javax.annotation.Nullable; import javax.net.ssl.HttpsURLConnection; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Discovers a remote service and resolves URLs to that service. @@ -59,6 +65,12 @@ public class RemoteClient { public static final String RUNTIME_SERVICE_ROUTING_BASE_URI = "cdap.runtime.service.routing.base.uri"; + private static final Logger LOG = LoggerFactory.getLogger(RemoteClient.class); + + private static final ThreadLocal CURRENT_RESOLVED_POD = new ThreadLocal<>(); + private static final ThreadLocal CURRENT_ROUTING_KEY = new ThreadLocal<>(); + private static final String TASK_MANAGER_URL = "http://cdap-task-manager.default.svc.cluster.local:11025"; + private static final Gson GSON = new Gson(); private final InternalAuthenticator internalAuthenticator; private final EndpointStrategy endpointStrategy; @@ -66,6 +78,7 @@ public class RemoteClient { private final String discoverableServiceName; private final String basePath; private final RemoteAuthenticator remoteAuthenticator; + private final DiscoveryServiceClient discoveryClient; RemoteClient(InternalAuthenticator internalAuthenticator, DiscoveryServiceClient discoveryClient, String discoverableServiceName, HttpRequestConfig httpRequestConfig, String basePath, @@ -73,6 +86,7 @@ public class RemoteClient { this.internalAuthenticator = internalAuthenticator; this.discoverableServiceName = discoverableServiceName; this.httpRequestConfig = httpRequestConfig; + this.discoveryClient = discoveryClient; this.endpointStrategy = new RandomEndpointStrategy( () -> discoveryClient.discover(discoverableServiceName)); String cleanBasePath = basePath.startsWith("/") ? basePath.substring(1) : basePath; @@ -93,6 +107,14 @@ public HttpRequest.Builder requestBuilder(HttpMethod method, String resource) { return HttpRequest.builder(method, resolve(resource)); } + /** + * Create a {@link HttpRequest.Builder} using the specified http method, resource, and routing key (namespace). + * This client will discover the service address and resolve it stickily using the routing key. + */ + public HttpRequest.Builder requestBuilder(HttpMethod method, String resource, @Nullable String routingKey) { + return HttpRequest.builder(method, resolve(resource, routingKey)); + } + private void setAuthHeader(BiConsumer headerSetter, String header, String credentialType, String credentialValue) { @@ -190,6 +212,14 @@ private HttpResponse executeNonIdempotent(HttpRequest request) throws IOExceptio return response; } catch (ConnectException e) { throw new ServiceUnavailableException(discoverableServiceName, e); + } finally { + Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); + String routingKey = CURRENT_ROUTING_KEY.get(); + if (resolvedPod != null && routingKey != null) { + notifyTaskManagerFinished(routingKey, resolvedPod); + } + CURRENT_RESOLVED_POD.remove(); + CURRENT_ROUTING_KEY.remove(); } } @@ -204,14 +234,50 @@ public void executeStreamingRequest(HttpRequest request) HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength(), request.getConsumer()); - HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); + try { + HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); - if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { - throw new IOException( - String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), - httpResponse.getResponseCode())); + if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { + throw new IOException( + String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), + httpResponse.getResponseCode())); + } + httpResponse.consumeContent(); + } finally { + Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); + String routingKey = CURRENT_ROUTING_KEY.get(); + if (resolvedPod != null && routingKey != null) { + notifyTaskManagerFinished(routingKey, resolvedPod); + } + CURRENT_RESOLVED_POD.remove(); + CURRENT_ROUTING_KEY.remove(); + } + } + + private void notifyTaskManagerFinished(String namespace, Discoverable pod) { + try { + URL url = new URL(TASK_MANAGER_URL + "/v3/taskmanager/finish"); + TaskManagerHttpHandler.FinishRequest finishRequest = new TaskManagerHttpHandler.FinishRequest(); + + java.lang.reflect.Field nsField = finishRequest.getClass().getDeclaredField("namespace"); + nsField.setAccessible(true); + nsField.set(finishRequest, namespace); + + TaskManagerHttpHandler.PodInfo podInfo = new TaskManagerHttpHandler.PodInfo( + pod.getSocketAddress().getHostString(), pod.getSocketAddress().getPort()); + java.lang.reflect.Field podField = finishRequest.getClass().getDeclaredField("pod"); + podField.setAccessible(true); + podField.set(finishRequest, podInfo); + + HttpRequest req = HttpRequest.post(url) + .addHeader(HttpHeaders.CONTENT_TYPE, "application/json") + .withBody(GSON.toJson(finishRequest)) + .build(); + + HttpRequests.execute(req, httpRequestConfig); + } catch (Exception e) { + LOG.warn("sidhdirenge - Failed to notify Task Manager of task completion", e); } - httpResponse.consumeContent(); } /** @@ -260,16 +326,114 @@ public HttpURLConnection openConnection(HttpMethod method, String resource) thro * @throws ServiceUnavailableException if the service could not be discovered */ public URL resolve(String resource) { - Discoverable discoverable = endpointStrategy.pick(1L, TimeUnit.SECONDS); - if (discoverable == null) { + return resolve(resource, null); + } + + /** + * Discover the service address, then append the base path and specified resource to get the URL, + * using a routing key (e.g. namespace) to ensure sticky routing to the same pod. If routingKey is + * null, it falls back to the default random discovery strategy. + */ + public URL resolve(String resource, @Nullable String routingKey) { + if (routingKey == null) { + Discoverable discoverable = endpointStrategy.pick(1L, TimeUnit.SECONDS); + if (discoverable == null) { + throw new ServiceUnavailableException(discoverableServiceName); + } + URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); + try { + return rewriteUrl(uri.toURL()); + } catch (MalformedURLException e) { + throw new IllegalStateException( + String.format("Discovered service %s, but it announced malformed URL %s", + discoverableServiceName, uri), e); + } + } + + LOG.info("sidhdirenge - RemoteClient resolving stickily via TaskManager for service {} with routingKey: {}", + discoverableServiceName, routingKey); + + // 1. Fetch all currently discovered endpoints + Iterable discoverables = () -> discoveryClient.discover(discoverableServiceName) + .iterator(); + List list = new ArrayList<>(); + for (Discoverable d : discoverables) { + // Perform DNS lookup to resolve the service hostname into individual pod IPs (for headless services) + try { + java.net.InetAddress[] addresses = java.net.InetAddress.getAllByName( + d.getSocketAddress().getHostName()); + for (java.net.InetAddress addr : addresses) { + list.add(new Discoverable(d.getName(), + new java.net.InetSocketAddress(addr.getHostAddress(), d.getSocketAddress().getPort()), + d.getPayload())); + } + } catch (java.net.UnknownHostException e) { + // Fallback to original discoverable if DNS lookup fails + list.add(d); + } + } + + if (list.isEmpty()) { throw new ServiceUnavailableException(discoverableServiceName); } + // 2. Sort endpoints by IP address and port to ensure consistent ordering across all client instances + list.sort(Comparator.comparing((Discoverable d) -> d.getSocketAddress().getHostName()) + .thenComparingInt(d -> d.getSocketAddress().getPort())); + + // 3. Delegate to the standalone TaskManager Service over HTTP + Discoverable discoverable = null; + try { + URL url = new URL(TASK_MANAGER_URL + "/v3/taskmanager/resolve"); + TaskManagerHttpHandler.ResolveRequest resolveRequest = new TaskManagerHttpHandler.ResolveRequest(); + + java.lang.reflect.Field nsField = resolveRequest.getClass().getDeclaredField("namespace"); + nsField.setAccessible(true); + nsField.set(resolveRequest, routingKey); + + List podInfos = new ArrayList<>(); + for (Discoverable pod : list) { + podInfos.add(new TaskManagerHttpHandler.PodInfo( + pod.getSocketAddress().getHostString(), pod.getSocketAddress().getPort())); + } + java.lang.reflect.Field podsField = resolveRequest.getClass().getDeclaredField("pods"); + podsField.setAccessible(true); + podsField.set(resolveRequest, podInfos); + + HttpRequest req = HttpRequest.post(url) + .addHeader(HttpHeaders.CONTENT_TYPE, "application/json") + .withBody(GSON.toJson(resolveRequest)) + .build(); + + HttpResponse resp = HttpRequests.execute(req, httpRequestConfig); + if (resp.getResponseCode() == HttpURLConnection.HTTP_OK) { + TaskManagerHttpHandler.PodInfo selectedPodInfo = GSON.fromJson( + resp.getResponseBodyAsString(), TaskManagerHttpHandler.PodInfo.class); + discoverable = new Discoverable("task.worker", + new java.net.InetSocketAddress(selectedPodInfo.getHost(), selectedPodInfo.getPort())); + } + } catch (Exception e) { + LOG.warn("sidhdirenge - Failed to resolve pod via Task Manager HTTP Service. Falling back to local hashing.", e); + } + + // Fallback: If Task Manager is down or returns error, use standard consistent hashing + if (discoverable == null) { + int baseIndex = (routingKey.hashCode() & Integer.MAX_VALUE) % list.size(); + discoverable = list.get(baseIndex); + LOG.warn("sidhdirenge - TaskManager resolution failed. Falling back to default index {}", baseIndex); + } + + // Store resolved pod context in ThreadLocal for task execution callbacks + CURRENT_RESOLVED_POD.set(discoverable); + CURRENT_ROUTING_KEY.set(routingKey); + + LOG.info("sidhdirenge - Centralized TaskManager selected warm pod IP {} for routingKey: {}", + discoverable.getSocketAddress(), routingKey); + URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); try { return rewriteUrl(uri.toURL()); } catch (MalformedURLException e) { - // shouldn't happen. If it does, it means there is some bug in the service announcer throw new IllegalStateException( String.format("Discovered service %s, but it announced malformed URL %s", discoverableServiceName, uri), e); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java new file mode 100644 index 000000000000..7dc2204d2e42 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java @@ -0,0 +1,185 @@ +/* + * 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 org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.locks.ReentrantLock; +import javax.annotation.Nullable; +import org.apache.twill.discovery.Discoverable; + +/** + * Centralized Task Manager for orchestrating Warm Sticky Leases on Task Worker pods. + * This class coordinates leases, concurrency, and logical resets. + */ +public class TaskManager { + + private static final Logger LOG = LoggerFactory.getLogger(TaskManager.class); + private static final TaskManager INSTANCE = new TaskManager(); + + // Concurrency and task limits based on the design doc + private static final int MAX_CONCURRENT_TASKS_PER_POD = 10; + private static final int MAX_TOTAL_TASKS_BEFORE_RESET = 10; + + private final ReentrantLock lock = new ReentrantLock(); + + // Lease state maps + // Pod IP/Key -> Namespace currently leased + private final Map podLeases = new HashMap<>(); + // Pod IP/Key -> Active concurrent task count + private final Map podActiveTaskCounts = new HashMap<>(); + // Pod IP/Key -> Total tasks processed on the current lease + private final Map podTotalTaskProcessedCounts = new HashMap<>(); + + public static TaskManager getInstance() { + return INSTANCE; + } + + private TaskManager() { + // Singleton + } + + /** + * Resolves the target warm pod for a given namespace based on the sticky lease model. + * + * @param namespace the namespace requesting execution + * @param availablePods the list of currently discovered pods + * @return the selected pod, or null if no pod is available + */ + @Nullable + public Discoverable resolvePod(String namespace, List availablePods) { + lock.lock(); + try { + String leasedPodIp = null; + Discoverable selectedPod = null; + + // 1. Find if a pod is already leased to this namespace and has capacity + for (Discoverable pod : availablePods) { + String podIp = getPodKey(pod); + String currentLease = podLeases.get(podIp); + + if (namespace.equals(currentLease)) { + int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); + if (activeTasks < MAX_CONCURRENT_TASKS_PER_POD) { + leasedPodIp = podIp; + selectedPod = pod; + break; + } + } + } + + // 2. If no active lease exists (or it is at capacity), find an idle/unleased pod + if (selectedPod == null) { + for (Discoverable pod : availablePods) { + String podIp = getPodKey(pod); + String currentLease = podLeases.get(podIp); + + if (currentLease == null) { + // Establish a new lease on this idle pod + podLeases.put(podIp, namespace); + podActiveTaskCounts.put(podIp, 0); + podTotalTaskProcessedCounts.put(podIp, 0); + + LOG.info("sidhdirenge - TaskManager: Established new lease for namespace '{}' on pod '{}'", + namespace, podIp); + + leasedPodIp = podIp; + selectedPod = pod; + break; + } + } + } + + // 3. Fallback: If all pods are leased to other namespaces, find the pod with the least load + if (selectedPod == null) { + LOG.warn("sidhdirenge - TaskManager: All pods are leased. Falling back to least-loaded pod."); + int minLoad = Integer.MAX_VALUE; + for (Discoverable pod : availablePods) { + String podIp = getPodKey(pod); + int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); + if (activeTasks < minLoad) { + minLoad = activeTasks; + selectedPod = pod; + leasedPodIp = podIp; + } + } + + // Force-assign lease to the new namespace + if (selectedPod != null) { + podLeases.put(leasedPodIp, namespace); + podActiveTaskCounts.put(leasedPodIp, 0); + podTotalTaskProcessedCounts.put(leasedPodIp, 0); + } + } + + // 4. Increment task counts for the selected pod + if (selectedPod != null) { + int activeTasks = podActiveTaskCounts.getOrDefault(leasedPodIp, 0) + 1; + int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(leasedPodIp, 0) + 1; + + podActiveTaskCounts.put(leasedPodIp, activeTasks); + podTotalTaskProcessedCounts.put(leasedPodIp, totalProcessed); + + LOG.info("sidhdirenge - TaskManager: Routing task for '{}' to pod '{}' (Active: {}, Total: {})", + namespace, leasedPodIp, activeTasks, totalProcessed); + + // Check if the pod has reached the logical reset threshold + if (totalProcessed >= MAX_TOTAL_TASKS_BEFORE_RESET) { + LOG.info("sidhdirenge - TaskManager: Pod '{}' reached logical reset threshold ({} tasks). Reclaiming lease.", + leasedPodIp, totalProcessed); + releaseLease(leasedPodIp); + } + } + + return selectedPod; + } finally { + lock.unlock(); + } + } + + /** + * Called when a task completes to decrement the active task count on the pod. + */ + public void finishTask(String namespace, Discoverable pod) { + lock.lock(); + try { + String podIp = getPodKey(pod); + int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); + if (activeTasks > 0) { + podActiveTaskCounts.put(podIp, activeTasks - 1); + } + LOG.info("sidhdirenge - TaskManager: Task finished for '{}' on pod '{}' (Remaining active: {})", + namespace, podIp, podActiveTaskCounts.getOrDefault(podIp, 0)); + } finally { + lock.unlock(); + } + } + + private void releaseLease(String podIp) { + podLeases.remove(podIp); + podActiveTaskCounts.remove(podIp); + podTotalTaskProcessedCounts.remove(podIp); + } + + private String getPodKey(Discoverable pod) { + return pod.getSocketAddress().getHostString() + ":" + pod.getSocketAddress().getPort(); + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java new file mode 100644 index 000000000000..a8818daaced2 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java @@ -0,0 +1,152 @@ +/* + * 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 com.google.gson.Gson; +import com.google.gson.reflect.TypeToken; +import io.cdap.http.AbstractHttpHandler; +import io.cdap.http.HttpResponder; +import io.netty.handler.codec.http.FullHttpRequest; +import io.netty.handler.codec.http.HttpResponseStatus; +import org.apache.twill.discovery.Discoverable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.lang.reflect.Type; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import javax.ws.rs.POST; +import javax.ws.rs.Path; + +/** + * Netty HTTP Handler for the standalone Task Manager Service. + */ +@Path("/v3/taskmanager") +public class TaskManagerHttpHandler extends AbstractHttpHandler { + + private static final Logger LOG = LoggerFactory.getLogger(TaskManagerHttpHandler.class); + private static final Gson GSON = new Gson(); + private final TaskManager taskManager = TaskManager.getInstance(); + + @POST + @Path("/resolve") + public void resolve(FullHttpRequest request, HttpResponder responder) { + try { + String jsonBody = request.content().toString(StandardCharsets.UTF_8); + ResolveRequest resolveRequest = GSON.fromJson(jsonBody, ResolveRequest.class); + + if (resolveRequest == null || resolveRequest.getNamespace() == null || resolveRequest.getPods() == null) { + responder.sendStatus(HttpResponseStatus.BAD_REQUEST); + return; + } + + // Convert serialized pods back to Discoverable objects + List discoverables = new ArrayList<>(); + for (PodInfo podInfo : resolveRequest.getPods()) { + discoverables.add(new Discoverable("task.worker", + new InetSocketAddress(podInfo.getHost(), podInfo.getPort()))); + } + + Discoverable selectedPod = taskManager.resolvePod(resolveRequest.getNamespace(), discoverables); + + if (selectedPod == null) { + responder.sendStatus(HttpResponseStatus.SERVICE_UNAVAILABLE); + return; + } + + PodInfo responsePod = new PodInfo( + selectedPod.getSocketAddress().getHostString(), + selectedPod.getSocketAddress().getPort() + ); + + responder.sendJson(HttpResponseStatus.OK, GSON.toJson(responsePod)); + } catch (Exception e) { + LOG.error("Failed to resolve pod in Task Manager Service", e); + responder.sendString(HttpResponseStatus.INTERNAL_SERVER_ERROR, e.getMessage()); + } + } + + @POST + @Path("/finish") + public void finish(FullHttpRequest request, HttpResponder responder) { + try { + String jsonBody = request.content().toString(StandardCharsets.UTF_8); + FinishRequest finishRequest = GSON.fromJson(jsonBody, FinishRequest.class); + + if (finishRequest == null || finishRequest.getNamespace() == null || finishRequest.getPod() == null) { + responder.sendStatus(HttpResponseStatus.BAD_REQUEST); + return; + } + + Discoverable pod = new Discoverable("task.worker", + new InetSocketAddress(finishRequest.getPod().getHost(), finishRequest.getPod().getPort())); + + taskManager.finishTask(finishRequest.getNamespace(), pod); + responder.sendStatus(HttpResponseStatus.OK); + } catch (Exception e) { + LOG.error("Failed to finish task in Task Manager Service", e); + responder.sendString(HttpResponseStatus.INTERNAL_SERVER_ERROR, e.getMessage()); + } + } + + // DTO Classes for Serialization + public static class ResolveRequest { + private String namespace; + private List pods; + + public String getNamespace() { + return namespace; + } + + public List getPods() { + return pods; + } + } + + public static class FinishRequest { + private String namespace; + private PodInfo pod; + + public String getNamespace() { + return namespace; + } + + public PodInfo getPod() { + return pod; + } + } + + public static class PodInfo { + private String host; + private int port; + + public PodInfo(String host, int port) { + this.host = host; + this.port = port; + } + + public String getHost() { + return host; + } + + public int getPort() { + return port; + } + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java new file mode 100644 index 000000000000..282d005d719e --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java @@ -0,0 +1,55 @@ +/* + * 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.cdap.http.NettyHttpService; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Main class for running the standalone Task Manager Service. + */ +public class TaskManagerMain { + + private static final Logger LOG = LoggerFactory.getLogger(TaskManagerMain.class); + private static final int PORT = 11025; + + public static void main(String[] args) throws Exception { + LOG.info("Starting standalone CDAP Task Manager Service on port {}...", PORT); + + NettyHttpService httpService = NettyHttpService.builder("task-manager") + .setHttpHandlers(new TaskManagerHttpHandler()) + .setPort(PORT) + .build(); + + httpService.start(); + LOG.info("CDAP Task Manager Service started successfully."); + + // Keep the service running + Runtime.getRuntime().addShutdownHook(new Thread(() -> { + try { + LOG.info("Stopping CDAP Task Manager Service..."); + httpService.stop(); + LOG.info("CDAP Task Manager Service stopped."); + } catch (Exception e) { + LOG.error("Error stopping Task Manager Service", e); + } + })); + + Thread.currentThread().join(); + } +} diff --git a/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java b/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java new file mode 100644 index 000000000000..acd746293720 --- /dev/null +++ b/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java @@ -0,0 +1,90 @@ +/* + * 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 org.apache.twill.discovery.Discoverable; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import java.net.InetSocketAddress; +import java.util.ArrayList; +import java.util.List; + +/** + * Unit tests for {@link TaskManager} warm sticky lease orchestration. + */ +public class TaskManagerTest { + + private List pods; + + @Before + public void setUp() { + pods = new ArrayList<>(); + // Define 3 mock task worker pods + pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.1", 11015))); + pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.2", 11015))); + pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.3", 11015))); + } + + @Test + public void testStickyRoutingAndLeasing() { + TaskManager taskManager = TaskManager.getInstance(); + + // 1. First request for ns1 should claim a pod + Discoverable pod1 = taskManager.resolvePod("ns1", pods); + Assert.assertNotNull(pod1); + + // 2. Second request for ns1 should land on the same pod (Stickiness) + Discoverable pod1Repeat = taskManager.resolvePod("ns1", pods); + Assert.assertEquals(pod1.getSocketAddress(), pod1Repeat.getSocketAddress()); + + // 3. First request for ns2 should claim a different, idle pod (Isolation) + Discoverable pod2 = taskManager.resolvePod("ns2", pods); + Assert.assertNotNull(pod2); + Assert.assertNotEquals(pod1.getSocketAddress(), pod2.getSocketAddress()); + + // Clean up active tasks + taskManager.finishTask("ns1", pod1); + taskManager.finishTask("ns1", pod1Repeat); + taskManager.finishTask("ns2", pod2); + } + + @Test + public void testLogicalResetAfterMaxTasks() { + TaskManager taskManager = TaskManager.getInstance(); + + // 1. Claim a pod for ns3 + Discoverable initialPod = taskManager.resolvePod("ns3", pods); + Assert.assertNotNull(initialPod); + taskManager.finishTask("ns3", initialPod); + + // 2. Send 9 more tasks to reach the limit of 10 total tasks + for (int i = 0; i < 9; i++) { + Discoverable p = taskManager.resolvePod("ns3", pods); + Assert.assertEquals(initialPod.getSocketAddress(), p.getSocketAddress()); + taskManager.finishTask("ns3", p); + } + + // 3. The 11th request for a DIFFERENT namespace (ns4) should now be able to claim this pod + // because it was logically reset (released) on the 10th task! + Discoverable resetPod = taskManager.resolvePod("ns4", pods); + Assert.assertNotNull(resetPod); + Assert.assertEquals(initialPod.getSocketAddress(), resetPod.getSocketAddress()); + taskManager.finishTask("ns4", resetPod); + } +} diff --git a/task-manager-service.yaml b/task-manager-service.yaml new file mode 100644 index 000000000000..ed4da21222b7 --- /dev/null +++ b/task-manager-service.yaml @@ -0,0 +1,49 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: cdap-task-manager + namespace: default + labels: + cdap.service: task.manager +spec: + replicas: 1 + selector: + matchLabels: + cdap.service: task.manager + template: + metadata: + labels: + cdap.service: task.manager + spec: + containers: + - name: task-manager + image: us-east1-docker.pkg.dev/ld27be8c949817660-tp/ar-demo/cloud-data-fusion:latest + command: + - "java" + - "-cp" + - "/opt/cdap/master/lib/*" + - "io.cdap.cdap.common.internal.remote.TaskManagerMain" + ports: + - containerPort: 11025 + name: http + resources: + requests: + cpu: "500m" + memory: "1024Mi" + limits: + cpu: "1000m" + memory: "2048Mi" +--- +apiVersion: v1 +kind: Service +metadata: + name: cdap-task-manager + namespace: default +spec: + selector: + cdap.service: task.manager + ports: + - protocol: TCP + port: 11025 + targetPort: 11025 + type: ClusterIP From 03d9761de1f9001a8a79fae52e2106da4ddfbd5a Mon Sep 17 00:00:00 2001 From: sidhdirenge Date: Mon, 29 Jun 2026 10:41:55 +0000 Subject: [PATCH 02/54] feat(rbac-poc): pass namespace as routingKey in RemoteTaskExecutor --- .../cdap/cdap/common/internal/remote/RemoteTaskExecutor.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index bb1f11fd1f34..3004f4f064ce 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -130,8 +130,9 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception try { return Retries.callWithRetries((retryContext) -> { try { + String namespace = runnableTaskRequest.getNamespace(); HttpRequest.Builder requestBuilder = remoteClient - .requestBuilder(HttpMethod.POST, workerUrl) + .requestBuilder(HttpMethod.POST, workerUrl, namespace) .withBody(requestBody.duplicate()); if (compression) { requestBuilder.addHeader(HttpHeaders.CONTENT_ENCODING, "gzip"); From 94f90793f41ec994ca969d77fb90abc36927bb93 Mon Sep 17 00:00:00 2001 From: sidhdirenge Date: Mon, 29 Jun 2026 10:50:52 +0000 Subject: [PATCH 03/54] feat(rbac-poc): integrate StickyLeaseManager on Task Worker side --- .../internal/remote/StickyLeaseManager.java | 199 ++++++++++++++++++ .../common/internal/remote/TaskDetails.java | 11 + .../remote/TaskWorkerHttpHandlerInternal.java | 198 +++++++++++------ 3 files changed, 339 insertions(+), 69 deletions(-) create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java new file mode 100644 index 000000000000..12754c57c642 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -0,0 +1,199 @@ +/* + * 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.cdap.cdap.proto.id.NamespaceId; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Proof of Concept (POC) Manager for Sticky Lease with Lifecycle & Reclamation. + *

+ * Demonstrates: 1. "First-Write Wins" Lease Acquisition (pinning pod to a namespace). 2. + * Enforcement: Rejection of mismatching namespace requests (429 TOO_MANY_REQUESTS). 3. Concurrency + * Control: Up to 10 concurrent tasks. 4. Logical Reset & Reclamation: - Task-Count Reclamation + * (releasing lease after 10 total tasks). - Tiered Inactivity Timeouts (Enterprise: 30s, Basic: + * 10s, Developer: 5s). 5. Switching Delay: Instant soft reset in < 10ms avoiding 40s cold boot + * penalty. + */ +public class StickyLeaseManager { + + private static final Logger LOG = LoggerFactory.getLogger(StickyLeaseManager.class); + private final int maxConcurrentTasks; + private final int maxTasksPerLease; + private final AtomicReference currentLease = new AtomicReference<>(null); + private final AtomicReference currentTier = new AtomicReference<>(TenantTier.BASIC); + private final AtomicInteger activeTaskCount = new AtomicInteger(0); + private final AtomicInteger totalTasksProcessedInLease = new AtomicInteger(0); + private volatile long lastActivityTimeMillis; + + public StickyLeaseManager() { + this(10, 10); + } + + public StickyLeaseManager(int maxConcurrentTasks, int maxTasksPerLease) { + this.maxConcurrentTasks = maxConcurrentTasks; + this.maxTasksPerLease = maxTasksPerLease; + this.lastActivityTimeMillis = System.currentTimeMillis(); + } + + /** + * Attempts to acquire or verify the sticky lease using "First-Write Wins" logic. + */ + public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, TenantTier tier) { + NamespaceId existing = currentLease.get(); + + if (existing == null) { + // Idle pod claimed by new namespace in < 10ms (Soft Reset / Claim) + long claimStartTime = System.currentTimeMillis(); + if (currentLease.compareAndSet(null, namespace)) { + currentTier.set(tier); + activeTaskCount.set(0); + totalTasksProcessedInLease.set(0); + lastActivityTimeMillis = System.currentTimeMillis(); + long elapsed = System.currentTimeMillis() - claimStartTime; + LOG.info( + "Lease claimed by namespace '{}' (Tier: {}) in {}ms (Boot penalty entirely avoided)", + namespace.getNamespace(), tier, elapsed); + return AcquisitionStatus.SUCCESS; + } + } + + // Check if matching existing lease + if (namespace.equals(currentLease.get())) { + lastActivityTimeMillis = System.currentTimeMillis(); + return AcquisitionStatus.SUCCESS; + } + + // Mismatching namespace -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) + LOG.info("Enforcement: Rejecting request for namespace '{}', current lease is held by '{}'", + namespace.getNamespace(), currentLease.get()); + return AcquisitionStatus.REJECTED_MISMATCH; + } + + /** + * Starts a task for the given namespace if lease and concurrency allow. + */ + public synchronized AcquisitionStatus startTask(NamespaceId namespace, TenantTier tier) { + AcquisitionStatus status = acquireLease(namespace, tier); + if (status != AcquisitionStatus.SUCCESS) { + return status; + } + + if (activeTaskCount.get() >= maxConcurrentTasks) { + LOG.info("Concurrency limit reached ({} tasks active) for namespace '{}'", + activeTaskCount.get(), namespace.getNamespace()); + return AcquisitionStatus.REJECTED_MAX_CONCURRENCY; + } + + activeTaskCount.incrementAndGet(); + lastActivityTimeMillis = System.currentTimeMillis(); + return AcquisitionStatus.SUCCESS; + } + + /** + * Finishes a task and evaluates Reclamation / Logical Reset after 10 total tasks. + */ + public synchronized void finishTask(NamespaceId namespace) { + if (namespace.equals(currentLease.get())) { + activeTaskCount.decrementAndGet(); + lastActivityTimeMillis = System.currentTimeMillis(); + int completed = totalTasksProcessedInLease.incrementAndGet(); + + // Reclamation: After processing 10 total tasks, release lease (Logical Reset) + if (completed >= maxTasksPerLease && activeTaskCount.get() == 0) { + releaseLease( + "Processed " + completed + " total tasks (Max " + maxTasksPerLease + " reached)"); + } + } + } + + /** + * Checks if the idle timeout for the current tiered tenancy has been exceeded. If exceeded, + * triggers a logical reset. + */ + public synchronized boolean enforceInactivityReclamation() { + NamespaceId leased = currentLease.get(); + if (leased != null && activeTaskCount.get() == 0) { + long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; + long threshold = currentTier.get().getInactivityTimeoutMillis(); + + if (idleDurationMillis >= threshold) { + releaseLease(String.format("Tiered inactivity timeout exceeded for %s (%dms >= %dms)", + currentTier.get(), idleDurationMillis, threshold)); + return true; + } + } + return false; + } + + /** + * Releases the lease (Logical Reset), clearing internal namespace context and sidecar state. + */ + public synchronized void releaseLease(String reason) { + NamespaceId oldNamespace = currentLease.getAndSet(null); + if (oldNamespace != null) { + activeTaskCount.set(0); + totalTasksProcessedInLease.set(0); + LOG.info("Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", + oldNamespace.getNamespace(), reason); + } + } + + @Nullable + public NamespaceId getCurrentLease() { + return currentLease.get(); + } + + public int getActiveTaskCount() { + return activeTaskCount.get(); + } + + public int getTotalTasksProcessedInLease() { + return totalTasksProcessedInLease.get(); + } + + public void setLastActivityTimeMillis(long timestampMillis) { + this.lastActivityTimeMillis = timestampMillis; + } + + public enum TenantTier { + ENTERPRISE(TimeUnit.SECONDS.toMillis(30)), + BASIC(TimeUnit.SECONDS.toMillis(10)), + DEVELOPER(TimeUnit.SECONDS.toMillis(5)); + + private final long inactivityTimeoutMillis; + + TenantTier(long inactivityTimeoutMillis) { + this.inactivityTimeoutMillis = inactivityTimeoutMillis; + } + + public long getInactivityTimeoutMillis() { + return inactivityTimeoutMillis; + } + } + + public enum AcquisitionStatus { + SUCCESS, + REJECTED_MISMATCH, + REJECTED_MAX_CONCURRENCY + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskDetails.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskDetails.java index d1e96f37e21a..19e1c610c24b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskDetails.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskDetails.java @@ -72,4 +72,15 @@ public String getClassName() { .map(RunnableTaskRequest::getClassName) .orElse(request.getClassName()); } + + @Nullable + public String getNamespace() { + if (request == null) { + return null; + } + return Optional.ofNullable(request.getParam()) + .map(RunnableTaskParam::getEmbeddedTaskRequest) + .map(RunnableTaskRequest::getNamespace) + .orElse(request.getNamespace()); + } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java index c2b2b52757cf..b7726c46af23 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java @@ -85,7 +85,8 @@ public class TaskWorkerHttpHandlerInternal extends AbstractHttpHandler { /** * Holds the total number of requests that have been executed by this handler - * that should count toward max allowed. + * that should count + * toward max allowed. */ private final AtomicInteger runningRequestCount = new AtomicInteger(0); private final AtomicInteger requestProcessedCount = new AtomicInteger(0); @@ -99,6 +100,7 @@ public class TaskWorkerHttpHandlerInternal extends AbstractHttpHandler { */ private final AtomicBoolean mustRestart = new AtomicBoolean(false); private final int concurrentRequestLimit; + private final StickyLeaseManager stickyLeaseManager; /** * Constructs the {@link TaskWorkerHttpHandlerInternal}. @@ -117,12 +119,19 @@ public TaskWorkerHttpHandlerInternal(CConfiguration cConf, Constants.TaskWorker.METADATA_SERVICE_END_POINT); boolean enableUserCodeIsolationEnabled = cConf.getBoolean( TaskWorker.USER_CODE_ISOLATION_ENABLED); - if (enableUserCodeIsolationEnabled) { - // Run only one request at a time in user code isolation mode. - this.concurrentRequestLimit = 1; - } else { - this.concurrentRequestLimit = cConf.getInt(TaskWorker.REQUEST_LIMIT); - } + this.concurrentRequestLimit = cConf.getInt(TaskWorker.REQUEST_LIMIT); + int maxTasksPerLease = cConf.getInt("task.worker.lease.max.tasks", 10); + this.stickyLeaseManager = new StickyLeaseManager(concurrentRequestLimit, maxTasksPerLease); + + ScheduledExecutorService leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( + Threads.createDaemonThreadFactory("lease-reclamation")); + leaseReclamationExecutor.scheduleAtFixedRate(() -> { + try { + stickyLeaseManager.enforceInactivityReclamation(); + } catch (Throwable t) { + LOG.warn("Error enforcing inactivity lease reclamation", t); + } + }, 1, 1, TimeUnit.SECONDS); // Restart the service to clean up and re-claim resources after user code // execution. @@ -131,6 +140,11 @@ public TaskWorkerHttpHandlerInternal(CConfiguration cConf, final int pendingRequests = runningRequestCount.decrementAndGet(); requestProcessedCount.incrementAndGet(); + String namespace = taskDetails.getNamespace(); + if (namespace != null) { + stickyLeaseManager.finishTask(new NamespaceId(namespace)); + } + String className = taskDetails.getClassName(); if (mustRestart.get() && pendingRequests == 0) { stopper.accept(className); @@ -154,11 +168,14 @@ public TaskWorkerHttpHandlerInternal(CConfiguration cConf, /** * If there is no ongoing request, worker pod gets restarted after a random - * duration is selected from the following range. Otherwise, worker pod can - * only get restarted once the ongoing request finishes. range = [Duration - - * DURATION_FRACTION * Duration, Duration + DURATION_FRACTION * Duration] - * Reason: by randomizing the duration, it is guaranteed that pods do not get - * restarted at the same time. + * duration is selected + * from the following range. Otherwise, worker pod can only get restarted once + * the ongoing request + * finishes. range = [Duration - DURATION_FRACTION * Duration, Duration + + * DURATION_FRACTION * + * Duration] Reason: by randomizing the duration, it is guaranteed that pods do + * not get restarted + * at the same time. */ private void enablePeriodicRestart(CConfiguration cConf, Consumer stopper) { @@ -174,7 +191,7 @@ private void enablePeriodicRestart(CConfiguration cConf, int finalTaskDeadlineSeconds = calculateFinalTaskDeadlineSeconds(duration); ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor( - Threads.createDaemonThreadFactory("task-worker-restart")); + Threads.createDaemonThreadFactory("task-worker-restart")); executorService.scheduleWithFixedDelay(() -> { // we restart once all ongoing requests finish, i.e. runningRequestCount is 0. @@ -211,71 +228,110 @@ private void stopAndShutdown(ScheduledExecutorService executorService, Consumer< @POST @Path("/run") public void run(FullHttpRequest request, HttpResponder responder) { + LOG.info("sidhdirenge - Received task on worker {} for namespace :{}", + System.getenv("HOSTNAME") != null ? System.getenv("HOSTNAME") : "unknown", + request.headers()); if (mustRestart.get()) { responder.sendStatus(HttpResponseStatus.TOO_MANY_REQUESTS); return; } - if (runningRequestCount.incrementAndGet() > concurrentRequestLimit) { - responder.sendStatus(HttpResponseStatus.TOO_MANY_REQUESTS); - runningRequestCount.decrementAndGet(); - return; - } long startTime = System.currentTimeMillis(); + RunnableTaskRequest runnableTaskRequest = null; try { - RunnableTaskRequest runnableTaskRequest = GSON.fromJson( + runnableTaskRequest = GSON.fromJson( request.content().toString(StandardCharsets.UTF_8), RunnableTaskRequest.class); - RunnableTaskContext runnableTaskContext = new RunnableTaskContext( - runnableTaskRequest); - try { - NamespaceId namespaceId; - if (runnableTaskRequest.getParam().getEmbeddedTaskRequest() != null) { - // For system app tasks - namespaceId = new NamespaceId( - runnableTaskRequest.getParam().getEmbeddedTaskRequest() - .getNamespace()); - } else { - namespaceId = new NamespaceId(runnableTaskRequest.getNamespace()); - } - // set the GcpMetadataTaskContext before running the task. - GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, - cConf); - runnableTaskLauncher.launchRunnableTask(runnableTaskContext); - TaskDetails taskDetails = new TaskDetails(metricsCollectionService, - startTime, runnableTaskContext.isTerminateOnComplete(), - runnableTaskRequest); - responder.sendContent(HttpResponseStatus.OK, - new RunnableTaskBodyProducer(runnableTaskContext, - taskCompletionConsumer, taskDetails), - new DefaultHttpHeaders().add(HttpHeaders.CONTENT_TYPE, - MediaType.APPLICATION_OCTET_STREAM)); - } catch (ClassNotFoundException | ClassCastException ex) { - responder.sendString(HttpResponseStatus.BAD_REQUEST, - exceptionToJson(ex), - new DefaultHttpHeaders().set(HttpHeaders.CONTENT_TYPE, - "application/json")); - // Since the user class is not even loaded, no user code ran, hence it's ok to not terminate the runner - taskCompletionConsumer.accept(false, - new TaskDetails(metricsCollectionService, startTime, false, - runnableTaskRequest)); - } finally { - // clear the GcpMetadataTaskContext after the task is completed. - GcpMetadataTaskContextUtil.clearGcpMetadataTaskContext(cConf); - } } catch (Exception ex) { - LOG.error("Failed to run task {}", + LOG.error("Failed to parse task request {}", request.content().toString(StandardCharsets.UTF_8), ex); + responder.sendString(HttpResponseStatus.BAD_REQUEST, + exceptionToJson(ex), + new DefaultHttpHeaders().set(HttpHeaders.CONTENT_TYPE, "application/json")); + return; + } + + NamespaceId namespaceId; + if (runnableTaskRequest.getParam() != null + && runnableTaskRequest.getParam().getEmbeddedTaskRequest() != null) { + namespaceId = new NamespaceId( + runnableTaskRequest.getParam().getEmbeddedTaskRequest().getNamespace()); + } else { + String ns = runnableTaskRequest.getNamespace(); + namespaceId = new NamespaceId(ns != null ? ns : "default"); + } + + StickyLeaseManager.TenantTier tier = getTenantTier(namespaceId); + StickyLeaseManager.AcquisitionStatus leaseStatus = stickyLeaseManager.startTask(namespaceId, + tier); + if (leaseStatus != StickyLeaseManager.AcquisitionStatus.SUCCESS) { + LOG.warn("Rejecting request for namespace {} due to lease status: {}", namespaceId, + leaseStatus); + responder.sendStatus(HttpResponseStatus.TOO_MANY_REQUESTS); + return; + } + + runningRequestCount.incrementAndGet(); + + try { + // set the GcpMetadataTaskContext before running the task. + GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, cConf); + RunnableTaskContext runnableTaskContext = new RunnableTaskContext(runnableTaskRequest); + runnableTaskLauncher.launchRunnableTask(runnableTaskContext); + + TaskDetails taskDetails = new TaskDetails(metricsCollectionService, + startTime, runnableTaskContext.isTerminateOnComplete(), + runnableTaskRequest); + + responder.sendContent(HttpResponseStatus.OK, + new RunnableTaskBodyProducer(runnableTaskContext, + taskCompletionConsumer, taskDetails), + new DefaultHttpHeaders().add(HttpHeaders.CONTENT_TYPE, + MediaType.APPLICATION_OCTET_STREAM)); + } catch (ClassNotFoundException | ClassCastException ex) { + responder.sendString(HttpResponseStatus.BAD_REQUEST, + exceptionToJson(ex), + new DefaultHttpHeaders().set(HttpHeaders.CONTENT_TYPE, + "application/json")); + // Since the user class is not even loaded, no user code ran, hence it's ok to + // not terminate the runner + taskCompletionConsumer.accept(false, + new TaskDetails(metricsCollectionService, startTime, false, + runnableTaskRequest)); + } catch (Exception ex) { + LOG.error("Failed to run task {}", runnableTaskRequest, ex); responder.sendString(HttpResponseStatus.INTERNAL_SERVER_ERROR, exceptionToJson(ex), new DefaultHttpHeaders().set(HttpHeaders.CONTENT_TYPE, "application/json")); // Potentially ran user code, hence terminate the runner. taskCompletionConsumer.accept(false, - new TaskDetails(metricsCollectionService, startTime, true, null)); + new TaskDetails(metricsCollectionService, startTime, true, runnableTaskRequest)); + } finally { + // clear the GcpMetadataTaskContext after the task is completed. + try { + GcpMetadataTaskContextUtil.clearGcpMetadataTaskContext(cConf); + } catch (Exception e) { + LOG.warn("Failed to clear GCP metadata task context", e); + } } } + private StickyLeaseManager.TenantTier getTenantTier(NamespaceId namespaceId) { + String nsName = namespaceId.getNamespace().toLowerCase(); + if (nsName.contains("enterprise")) { + return StickyLeaseManager.TenantTier.ENTERPRISE; + } else if (nsName.contains("developer") || nsName.contains("dev")) { + return StickyLeaseManager.TenantTier.DEVELOPER; + } else { + return StickyLeaseManager.TenantTier.BASIC; + } + } + + StickyLeaseManager getStickyLeaseManager() { + return stickyLeaseManager; + } + /** * Returns a new token from metadata server. * @@ -309,7 +365,8 @@ public void token(io.netty.handler.codec.http.HttpRequest request, /** * Return json representation of an exception. Used to propagate exception - * across network for better surfacing errors and debuggability. + * across network for + * better surfacing errors and debuggability. */ private String exceptionToJson(Exception ex) { BasicThrowable basicThrowable = new BasicThrowable(ex); @@ -317,22 +374,24 @@ private String exceptionToJson(Exception ex) { } /** - * Compute the final task Dead line in Seconds where if the config {@TaskWorker.TASK_EXECUTION_DEADLINE_SECOND} - * is less than 0 which is not valid then use the duration instead. + * Compute the final task Dead line in Seconds where if the config + * {@TaskWorker.TASK_EXECUTION_DEADLINE_SECOND} is less than 0 which is not + * valid then use the + * duration instead. * * @param duration * @return */ private int calculateFinalTaskDeadlineSeconds(int duration) { int taskDeadlineSeconds = cConf.getInt( - TaskWorker.TASK_EXECUTION_DEADLINE_SECOND, - 0); + TaskWorker.TASK_EXECUTION_DEADLINE_SECOND, + 0); if (taskDeadlineSeconds < 0) { LOG.info( - "Task deadline is {}, using {} value {} as the deadline instead.", - taskDeadlineSeconds, - Constants.TaskWorker.CONTAINER_KILL_AFTER_DURATION_SECOND, duration); + "Task deadline is {}, using {} value {} as the deadline instead.", + taskDeadlineSeconds, + Constants.TaskWorker.CONTAINER_KILL_AFTER_DURATION_SECOND, duration); taskDeadlineSeconds = duration; } return taskDeadlineSeconds; @@ -340,9 +399,10 @@ private int calculateFinalTaskDeadlineSeconds(int duration) { /** * By using BodyProducer instead of simply sending out response bytes, the - * handler can get notified (through finished method) when sending the - * response is done, so it can safely call the stopper to kill the worker - * pod. + * handler can get + * notified (through finished method) when sending the response is done, so it + * can safely call the + * stopper to kill the worker pod. */ private static class RunnableTaskBodyProducer extends BodyProducer { From 375876bf16d63892e1becfb11e6813835678e406 Mon Sep 17 00:00:00 2001 From: sidhdirenge Date: Thu, 2 Jul 2026 09:19:00 +0000 Subject: [PATCH 04/54] feat(rbac-poc): refactor TaskManagerMain to extend AbstractServiceMain and use CommonNettyHttpServiceFactory --- .../internal/remote/TaskManagerMain.java | 55 -------------- .../internal/remote/TaskManagerService.java | 67 +++++++++++++++++ .../remote/TaskManagerServiceModule.java | 38 ++++++++++ .../environment/k8s/TaskManagerMain.java | 75 +++++++++++++++++++ task-manager-service.yaml | 57 ++++++++++++-- 5 files changed, 232 insertions(+), 60 deletions(-) delete mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java create mode 100644 cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java deleted file mode 100644 index 282d005d719e..000000000000 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerMain.java +++ /dev/null @@ -1,55 +0,0 @@ -/* - * 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.cdap.http.NettyHttpService; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Main class for running the standalone Task Manager Service. - */ -public class TaskManagerMain { - - private static final Logger LOG = LoggerFactory.getLogger(TaskManagerMain.class); - private static final int PORT = 11025; - - public static void main(String[] args) throws Exception { - LOG.info("Starting standalone CDAP Task Manager Service on port {}...", PORT); - - NettyHttpService httpService = NettyHttpService.builder("task-manager") - .setHttpHandlers(new TaskManagerHttpHandler()) - .setPort(PORT) - .build(); - - httpService.start(); - LOG.info("CDAP Task Manager Service started successfully."); - - // Keep the service running - Runtime.getRuntime().addShutdownHook(new Thread(() -> { - try { - LOG.info("Stopping CDAP Task Manager Service..."); - httpService.stop(); - LOG.info("CDAP Task Manager Service stopped."); - } catch (Exception e) { - LOG.error("Error stopping Task Manager Service", e); - } - })); - - Thread.currentThread().join(); - } -} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java new file mode 100644 index 000000000000..37510d9d7d3f --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -0,0 +1,67 @@ +/* + * 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 com.google.common.util.concurrent.AbstractIdleService; +import com.google.inject.Inject; +import io.cdap.cdap.common.conf.CConfiguration; +import io.cdap.cdap.common.http.CommonNettyHttpServiceFactory; +import io.cdap.http.NettyHttpService; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Collections; + +/** + * Guice-managed service that runs the Centralized Task Manager HTTP Server. + */ +public class TaskManagerService extends AbstractIdleService { + + private static final Logger LOG = LoggerFactory.getLogger(TaskManagerService.class); + private final NettyHttpService httpService; + + @Inject + TaskManagerService(CConfiguration cConf, + CommonNettyHttpServiceFactory commonNettyHttpServiceFactory, + TaskManagerHttpHandler taskManagerHttpHandler) { + + int port = cConf.getInt("task.manager.port", 11025); + String address = cConf.get("task.manager.address", "0.0.0.0"); + + LOG.info("sidhdirenge - Initializing TaskManagerService on {}:{}", address, port); + + this.httpService = commonNettyHttpServiceFactory.builder("task-manager", false) + .setHost(address) + .setPort(port) + .setHttpHandlers(Collections.singletonList(taskManagerHttpHandler)) + .build(); + } + + @Override + protected void startUp() throws Exception { + LOG.info("sidhdirenge - Starting TaskManagerService HTTP server..."); + httpService.start(); + LOG.info("sidhdirenge - TaskManagerService HTTP server started successfully at {}", httpService.getBindAddress()); + } + + @Override + protected void shutDown() throws Exception { + LOG.info("sidhdirenge - Stopping TaskManagerService HTTP server..."); + httpService.stop(); + LOG.info("sidhdirenge - TaskManagerService HTTP server stopped."); + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java new file mode 100644 index 000000000000..fb470462ab7e --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java @@ -0,0 +1,38 @@ +/* + * 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 com.google.inject.AbstractModule; +import com.google.inject.Scopes; + +/** + * Guice Module for Task Manager Service. + */ +public class TaskManagerServiceModule extends AbstractModule { + + @Override + protected void configure() { + // Bind the core TaskManager as a singleton + bind(TaskManager.class).toProvider(TaskManager::getInstance).in(Scopes.SINGLETON); + + // Bind the HTTP handler + bind(TaskManagerHttpHandler.class).in(Scopes.SINGLETON); + + // Bind the service itself + bind(TaskManagerService.class).in(Scopes.SINGLETON); + } +} diff --git a/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java new file mode 100644 index 000000000000..e71c70d877ff --- /dev/null +++ b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java @@ -0,0 +1,75 @@ +/* + * 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.master.environment.k8s; + +import com.google.common.util.concurrent.Service; +import com.google.inject.Injector; +import com.google.inject.Module; +import io.cdap.cdap.common.conf.CConfiguration; +import io.cdap.cdap.common.conf.Constants; +import io.cdap.cdap.common.internal.remote.TaskManagerService; +import io.cdap.cdap.common.internal.remote.TaskManagerServiceModule; +import io.cdap.cdap.common.logging.LoggingContext; +import io.cdap.cdap.common.logging.ServiceLoggingContext; +import io.cdap.cdap.master.spi.environment.MasterEnvironment; +import io.cdap.cdap.master.spi.environment.MasterEnvironmentContext; +import io.cdap.cdap.proto.id.NamespaceId; +import io.cdap.cdap.data2.audit.AuditModule; +import io.cdap.cdap.messaging.guice.MessagingServiceModule; + +import java.util.Arrays; +import java.util.List; +import javax.annotation.Nullable; + +/** + * Main entry point for the standalone Task Manager Service on Kubernetes. + */ +public class TaskManagerMain extends AbstractServiceMain { + + public static void main(String[] args) throws Exception { + main(TaskManagerMain.class, args); + } + + @Override + protected List getServiceModules(MasterEnvironment masterEnv, + EnvironmentOptions options, + CConfiguration cConf) { + return Arrays.asList( + new MessagingServiceModule(cConf), + new AuditModule(), + getDataFabricModule(), + new TaskManagerServiceModule() + ); + } + + @Override + protected void addServices(Injector injector, List services, + List closeableResources, + MasterEnvironment masterEnv, + MasterEnvironmentContext masterEnvContext, + EnvironmentOptions options) { + services.add(injector.getInstance(TaskManagerService.class)); + } + + @Nullable + @Override + protected LoggingContext getLoggingContext(EnvironmentOptions options) { + return new ServiceLoggingContext(NamespaceId.SYSTEM.getNamespace(), + Constants.Logging.COMPONENT_NAME, + "task-manager"); + } +} diff --git a/task-manager-service.yaml b/task-manager-service.yaml index ed4da21222b7..d14a24b60d35 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -5,6 +5,7 @@ metadata: namespace: default labels: cdap.service: task.manager + cdap.instance: sidhdirenge-jun8 spec: replicas: 1 selector: @@ -14,15 +15,19 @@ spec: metadata: labels: cdap.service: task.manager + cdap.instance: sidhdirenge-jun8 spec: + serviceAccountName: cdap-sidhdirenge-jun8-system-sa containers: - name: task-manager image: us-east1-docker.pkg.dev/ld27be8c949817660-tp/ar-demo/cloud-data-fusion:latest - command: - - "java" - - "-cp" - - "/opt/cdap/master/lib/*" - - "io.cdap.cdap.common.internal.remote.TaskManagerMain" + imagePullPolicy: Always + args: + - "io.cdap.cdap.master.environment.k8s.TaskManagerMain" + - "--env=k8s" + env: + - name: SERVICE_NAME + value: task-manager ports: - containerPort: 11025 name: http @@ -33,6 +38,48 @@ spec: limits: cpu: "1000m" memory: "2048Mi" + volumeMounts: + - mountPath: /etc/podinfo + name: podinfo + readOnly: true + - mountPath: /etc/cdap/conf + name: cdap-conf + readOnly: true + - mountPath: /etc/hadoop/conf + name: hadoop-conf + readOnly: true + - mountPath: /etc/cdap/security + name: cdap-security + readOnly: true + volumes: + - downwardAPI: + defaultMode: 420 + items: + - fieldRef: + apiVersion: v1 + fieldPath: metadata.labels + path: pod.labels.properties + - fieldRef: + apiVersion: v1 + fieldPath: metadata.name + path: pod.name + - fieldRef: + apiVersion: v1 + fieldPath: metadata.uid + path: pod.uid + name: podinfo + - configMap: + defaultMode: 420 + name: cdap-sidhdirenge-jun8-cconf + name: cdap-conf + - configMap: + defaultMode: 420 + name: cdap-sidhdirenge-jun8-hconf + name: hadoop-conf + - name: cdap-security + secret: + defaultMode: 420 + secretName: cdap-security --- apiVersion: v1 kind: Service From 5b162ec187616d1f8741f5e2acb4171aeac8cd51 Mon Sep 17 00:00:00 2001 From: sidhdirenge Date: Wed, 8 Jul 2026 07:07:59 +0000 Subject: [PATCH 05/54] feat(sticky-leases): implement randomized tie-breaking fallback and mismatch rejection count reversion in TaskManager --- .../common/internal/remote/RemoteClient.java | 39 ++++++++++-- .../internal/remote/RemoteTaskExecutor.java | 12 ++++ .../common/internal/remote/TaskManager.java | 60 ++++++++++++------- .../remote/TaskManagerHttpHandler.java | 7 ++- .../internal/remote/TaskManagerTest.java | 53 +++++++++++----- .../k8s/discovery/KubeDiscoveryService.java | 51 ++++++++++++++-- task-manager-service.yaml | 6 ++ 7 files changed, 179 insertions(+), 49 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index 7fb8bc8ee829..9bea26792ea0 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -181,9 +181,13 @@ private HttpResponse executeNonIdempotent(HttpRequest request) throws IOExceptio HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength()); + boolean rejected = false; try { HttpResponse response = HttpRequests.execute(httpRequest, httpRequestConfig); int responseCode = response.getResponseCode(); + if (responseCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { + rejected = true; + } // 503 is always retryable. Other 5xx errors are retryable if the request is idempotent (handled in // RemoteClient#executeIdempotent(HttpRequest) if (responseCode == HttpURLConnection.HTTP_UNAVAILABLE) { @@ -211,12 +215,16 @@ private HttpResponse executeNonIdempotent(HttpRequest request) throws IOExceptio } return response; } catch (ConnectException e) { + rejected = true; throw new ServiceUnavailableException(discoverableServiceName, e); + } catch (IOException | RuntimeException e) { + rejected = true; + throw e; } finally { Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); String routingKey = CURRENT_ROUTING_KEY.get(); if (resolvedPod != null && routingKey != null) { - notifyTaskManagerFinished(routingKey, resolvedPod); + notifyTaskManagerFinished(routingKey, resolvedPod, rejected); } CURRENT_RESOLVED_POD.remove(); CURRENT_ROUTING_KEY.remove(); @@ -234,27 +242,33 @@ public void executeStreamingRequest(HttpRequest request) HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength(), request.getConsumer()); + boolean rejected = false; try { HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { + if (httpResponse.getResponseCode() == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { + rejected = true; + } throw new IOException( String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), httpResponse.getResponseCode())); } httpResponse.consumeContent(); + } catch (IOException | RuntimeException e) { + rejected = true; + throw e; } finally { Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); String routingKey = CURRENT_ROUTING_KEY.get(); if (resolvedPod != null && routingKey != null) { - notifyTaskManagerFinished(routingKey, resolvedPod); + notifyTaskManagerFinished(routingKey, resolvedPod, rejected); } CURRENT_RESOLVED_POD.remove(); CURRENT_ROUTING_KEY.remove(); } } - - private void notifyTaskManagerFinished(String namespace, Discoverable pod) { + private void notifyTaskManagerFinished(String namespace, Discoverable pod, boolean rejected) { try { URL url = new URL(TASK_MANAGER_URL + "/v3/taskmanager/finish"); TaskManagerHttpHandler.FinishRequest finishRequest = new TaskManagerHttpHandler.FinishRequest(); @@ -269,6 +283,10 @@ private void notifyTaskManagerFinished(String namespace, Discoverable pod) { podField.setAccessible(true); podField.set(finishRequest, podInfo); + java.lang.reflect.Field rejectedField = finishRequest.getClass().getDeclaredField("rejected"); + rejectedField.setAccessible(true); + rejectedField.set(finishRequest, rejected); + HttpRequest req = HttpRequest.post(url) .addHeader(HttpHeaders.CONTENT_TYPE, "application/json") .withBody(GSON.toJson(finishRequest)) @@ -409,8 +427,19 @@ public URL resolve(String resource, @Nullable String routingKey) { if (resp.getResponseCode() == HttpURLConnection.HTTP_OK) { TaskManagerHttpHandler.PodInfo selectedPodInfo = GSON.fromJson( resp.getResponseBodyAsString(), TaskManagerHttpHandler.PodInfo.class); + + byte[] payload = list.isEmpty() ? new byte[0] : list.get(0).getPayload(); + for (Discoverable d : list) { + if (d.getSocketAddress().getPort() == selectedPodInfo.getPort() + && (d.getSocketAddress().getHostName().equals(selectedPodInfo.getHost()) + || (d.getSocketAddress().getAddress() != null + && d.getSocketAddress().getAddress().getHostAddress().equals(selectedPodInfo.getHost())))) { + payload = d.getPayload(); + break; + } + } discoverable = new Discoverable("task.worker", - new java.net.InetSocketAddress(selectedPodInfo.getHost(), selectedPodInfo.getPort())); + new java.net.InetSocketAddress(selectedPodInfo.getHost(), selectedPodInfo.getPort()), payload); } } catch (Exception e) { LOG.warn("sidhdirenge - Failed to resolve pod via Task Manager HTTP Service. Falling back to local hashing.", e); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 3004f4f064ce..aa4682c8fec2 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -43,6 +43,8 @@ import io.cdap.common.http.HttpRequestConfig; import io.cdap.common.http.HttpResponse; import io.netty.handler.codec.http.HttpResponseStatus; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.io.ByteArrayInputStream; import java.io.IOException; import java.io.InputStream; @@ -65,6 +67,7 @@ */ public class RemoteTaskExecutor { + private static final Logger LOG = LoggerFactory.getLogger(RemoteTaskExecutor.class); private static final Gson GSON = new Gson(); private static final String TASK_WORKER_URL = "/worker/run"; private static final String SYSTEM_WORKER_URL = "/system/run"; @@ -131,6 +134,15 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception return Retries.callWithRetries((retryContext) -> { try { String namespace = runnableTaskRequest.getNamespace(); + if ("system".equals(namespace) && runnableTaskRequest.getParam() != null + && runnableTaskRequest.getParam().getEmbeddedTaskRequest() != null) { + String embeddedNamespace = runnableTaskRequest.getParam().getEmbeddedTaskRequest().getNamespace(); + if (embeddedNamespace != null && !embeddedNamespace.isEmpty()) { + namespace = embeddedNamespace; + LOG.info("sidhdirenge - RemoteTaskExecutor: Mapped SystemAppTask namespace to embedded: {}", + namespace); + } + } HttpRequest.Builder requestBuilder = remoteClient .requestBuilder(HttpMethod.POST, workerUrl, namespace) .withBody(requestBody.duplicate()); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java index 7dc2204d2e42..0ab82fe9967f 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java @@ -16,15 +16,15 @@ package io.cdap.cdap.common.internal.remote; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.locks.ReentrantLock; import javax.annotation.Nullable; import org.apache.twill.discovery.Discoverable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Centralized Task Manager for orchestrating Warm Sticky Leases on Task Worker pods. @@ -37,7 +37,7 @@ public class TaskManager { // Concurrency and task limits based on the design doc private static final int MAX_CONCURRENT_TASKS_PER_POD = 10; - private static final int MAX_TOTAL_TASKS_BEFORE_RESET = 10; + static final int MAX_TOTAL_TASKS_BEFORE_RESET = 10; private final ReentrantLock lock = new ReentrantLock(); @@ -78,7 +78,8 @@ public Discoverable resolvePod(String namespace, List availablePod if (namespace.equals(currentLease)) { int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); - if (activeTasks < MAX_CONCURRENT_TASKS_PER_POD) { + int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); + if (activeTasks < MAX_CONCURRENT_TASKS_PER_POD && totalProcessed < MAX_TOTAL_TASKS_BEFORE_RESET) { leasedPodIp = podIp; selectedPod = pod; break; @@ -112,18 +113,25 @@ public Discoverable resolvePod(String namespace, List availablePod if (selectedPod == null) { LOG.warn("sidhdirenge - TaskManager: All pods are leased. Falling back to least-loaded pod."); int minLoad = Integer.MAX_VALUE; + List bestPods = new ArrayList<>(); for (Discoverable pod : availablePods) { String podIp = getPodKey(pod); int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); if (activeTasks < minLoad) { minLoad = activeTasks; - selectedPod = pod; - leasedPodIp = podIp; + bestPods.clear(); + bestPods.add(pod); + } else if (activeTasks == minLoad) { + bestPods.add(pod); } } - // Force-assign lease to the new namespace - if (selectedPod != null) { + if (!bestPods.isEmpty()) { + int randomIndex = java.util.concurrent.ThreadLocalRandom.current().nextInt(bestPods.size()); + selectedPod = bestPods.get(randomIndex); + leasedPodIp = getPodKey(selectedPod); + + // Force-assign lease to the new namespace podLeases.put(leasedPodIp, namespace); podActiveTaskCounts.put(leasedPodIp, 0); podTotalTaskProcessedCounts.put(leasedPodIp, 0); @@ -140,13 +148,6 @@ public Discoverable resolvePod(String namespace, List availablePod LOG.info("sidhdirenge - TaskManager: Routing task for '{}' to pod '{}' (Active: {}, Total: {})", namespace, leasedPodIp, activeTasks, totalProcessed); - - // Check if the pod has reached the logical reset threshold - if (totalProcessed >= MAX_TOTAL_TASKS_BEFORE_RESET) { - LOG.info("sidhdirenge - TaskManager: Pod '{}' reached logical reset threshold ({} tasks). Reclaiming lease.", - leasedPodIp, totalProcessed); - releaseLease(leasedPodIp); - } } return selectedPod; @@ -155,19 +156,32 @@ public Discoverable resolvePod(String namespace, List availablePod } } - /** - * Called when a task completes to decrement the active task count on the pod. - */ - public void finishTask(String namespace, Discoverable pod) { + public void finishTask(String namespace, Discoverable pod, boolean rejected) { lock.lock(); try { String podIp = getPodKey(pod); int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); if (activeTasks > 0) { - podActiveTaskCounts.put(podIp, activeTasks - 1); + activeTasks = activeTasks - 1; + podActiveTaskCounts.put(podIp, activeTasks); + } + + if (rejected) { + int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); + if (totalProcessed > 0) { + podTotalTaskProcessedCounts.put(podIp, totalProcessed - 1); + } + } else { + int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); + if (totalProcessed >= MAX_TOTAL_TASKS_BEFORE_RESET && activeTasks == 0) { + LOG.info("sidhdirenge - TaskManager: Pod '{}' finished all active tasks " + + "after reaching reset threshold. Reclaiming lease.", podIp); + releaseLease(podIp); + } } - LOG.info("sidhdirenge - TaskManager: Task finished for '{}' on pod '{}' (Remaining active: {})", - namespace, podIp, podActiveTaskCounts.getOrDefault(podIp, 0)); + + LOG.info("sidhdirenge - TaskManager: Task finished for '{}' on pod '{}' (Remaining active: {}, Rejected: {})", + namespace, podIp, activeTasks, rejected); } finally { lock.unlock(); } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java index a8818daaced2..d53ea147929b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java @@ -97,7 +97,7 @@ public void finish(FullHttpRequest request, HttpResponder responder) { Discoverable pod = new Discoverable("task.worker", new InetSocketAddress(finishRequest.getPod().getHost(), finishRequest.getPod().getPort())); - taskManager.finishTask(finishRequest.getNamespace(), pod); + taskManager.finishTask(finishRequest.getNamespace(), pod, finishRequest.isRejected()); responder.sendStatus(HttpResponseStatus.OK); } catch (Exception e) { LOG.error("Failed to finish task in Task Manager Service", e); @@ -122,6 +122,7 @@ public List getPods() { public static class FinishRequest { private String namespace; private PodInfo pod; + private boolean rejected; public String getNamespace() { return namespace; @@ -130,6 +131,10 @@ public String getNamespace() { public PodInfo getPod() { return pod; } + + public boolean isRejected() { + return rejected; + } } public static class PodInfo { diff --git a/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java b/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java index acd746293720..6ca5893c417c 100644 --- a/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java +++ b/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java @@ -16,15 +16,14 @@ package io.cdap.cdap.common.internal.remote; +import java.net.InetSocketAddress; +import java.util.ArrayList; +import java.util.List; import org.apache.twill.discovery.Discoverable; import org.junit.Assert; import org.junit.Before; import org.junit.Test; -import java.net.InetSocketAddress; -import java.util.ArrayList; -import java.util.List; - /** * Unit tests for {@link TaskManager} warm sticky lease orchestration. */ @@ -59,9 +58,9 @@ public void testStickyRoutingAndLeasing() { Assert.assertNotEquals(pod1.getSocketAddress(), pod2.getSocketAddress()); // Clean up active tasks - taskManager.finishTask("ns1", pod1); - taskManager.finishTask("ns1", pod1Repeat); - taskManager.finishTask("ns2", pod2); + taskManager.finishTask("ns1", pod1, false); + taskManager.finishTask("ns1", pod1Repeat, false); + taskManager.finishTask("ns2", pod2, false); } @Test @@ -71,20 +70,46 @@ public void testLogicalResetAfterMaxTasks() { // 1. Claim a pod for ns3 Discoverable initialPod = taskManager.resolvePod("ns3", pods); Assert.assertNotNull(initialPod); - taskManager.finishTask("ns3", initialPod); + taskManager.finishTask("ns3", initialPod, false); - // 2. Send 9 more tasks to reach the limit of 10 total tasks - for (int i = 0; i < 9; i++) { + // 2. Send tasks to reach the logical reset limit + for (int i = 0; i < TaskManager.MAX_TOTAL_TASKS_BEFORE_RESET - 1; i++) { Discoverable p = taskManager.resolvePod("ns3", pods); Assert.assertEquals(initialPod.getSocketAddress(), p.getSocketAddress()); - taskManager.finishTask("ns3", p); + taskManager.finishTask("ns3", p, false); } - // 3. The 11th request for a DIFFERENT namespace (ns4) should now be able to claim this pod - // because it was logically reset (released) on the 10th task! + // 3. The next request for a DIFFERENT namespace (ns4) should now be able to claim this pod + // because it was logically reset (released) on the limit threshold! Discoverable resetPod = taskManager.resolvePod("ns4", pods); Assert.assertNotNull(resetPod); Assert.assertEquals(initialPod.getSocketAddress(), resetPod.getSocketAddress()); - taskManager.finishTask("ns4", resetPod); + taskManager.finishTask("ns4", resetPod, false); + } + + @Test + public void testRejectionRevertsTotalProcessedCount() { + TaskManager taskManager = TaskManager.getInstance(); + + // 1. Claim a pod for ns5 + Discoverable initialPod = taskManager.resolvePod("ns5", pods); + Assert.assertNotNull(initialPod); + + // 2. Reject it -> totalProcessed should revert back to 0 + taskManager.finishTask("ns5", initialPod, true); + + // 3. Send 10 tasks to reach the logical reset limit + for (int i = 0; i < TaskManager.MAX_TOTAL_TASKS_BEFORE_RESET; i++) { + Discoverable p = taskManager.resolvePod("ns5", pods); + Assert.assertEquals(initialPod.getSocketAddress(), p.getSocketAddress()); + taskManager.finishTask("ns5", p, false); + } + + // 4. The next request for ns6 should claim this pod because it reset successfully after exactly 10 tasks, + // meaning the rejected task did not count towards the 10 task limit. + Discoverable resetPod = taskManager.resolvePod("ns6", pods); + Assert.assertNotNull(resetPod); + Assert.assertEquals(initialPod.getSocketAddress(), resetPod.getSocketAddress()); + taskManager.finishTask("ns6", resetPod, false); } } diff --git a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java index cd78998331bd..77d4b9ba0f22 100644 --- a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java +++ b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java @@ -23,6 +23,7 @@ import io.kubernetes.client.openapi.ApiClient; import io.kubernetes.client.openapi.ApiException; import io.kubernetes.client.openapi.apis.CoreV1Api; +import io.kubernetes.client.openapi.models.V1Endpoints; import io.kubernetes.client.openapi.models.V1ObjectMeta; import io.kubernetes.client.openapi.models.V1OwnerReference; import io.kubernetes.client.openapi.models.V1Pod; @@ -224,7 +225,7 @@ public void close() { * @throws IOException if exception was raised during creation of * {@link CoreV1Api} */ - private CoreV1Api getCoreApi() throws IOException { + CoreV1Api getCoreApi() throws IOException { CoreV1Api api = coreApi; if (api != null) { return api; @@ -542,7 +543,6 @@ Set toDiscoverables(String name, V1Service service, .map(Base64.getDecoder()::decode) .orElse(EMPTY_PAYLOAD); - String hostname; if (SERVICE_TYPE_LOAD_BALANCER.equals(service.getSpec().getType())) { Optional ipAddr = getLoadBalancerIp(service); if (!ipAddr.isPresent()) { @@ -551,12 +551,51 @@ Set toDiscoverables(String name, V1Service service, name); return Collections.emptySet(); } - hostname = ipAddr.get(); - } else { - hostname = String.format("%s.%s", meta.getName(), namespace); + String hostname = ipAddr.get(); + return servicePorts.stream() + .map(port -> createDiscoverable( + name, hostname, + port, payload) + ) + .filter(Objects::nonNull) + .findFirst() + .map(Collections::singleton) + .orElse(Collections.emptySet()); + } + + // Try to discover individual Pod IPs via the K8s Endpoints API for ClusterIP services + Set discoverables = new HashSet<>(); + try { + CoreV1Api api = getCoreApi(); + V1Endpoints endpoints = api.readNamespacedEndpoints(meta.getName(), namespace, null); + if (endpoints != null && endpoints.getSubsets() != null) { + for (io.kubernetes.client.openapi.models.V1EndpointSubset subset : endpoints.getSubsets()) { + List addresses = subset.getAddresses(); + List ports = subset.getPorts(); + if (addresses != null && ports != null) { + for (io.kubernetes.client.openapi.models.V1EndpointAddress address : addresses) { + for (io.kubernetes.client.openapi.models.CoreV1EndpointPort port : ports) { + if (servicePorts.stream().anyMatch(sp -> sp.getPort().equals(port.getPort()))) { + Discoverable d = createDiscoverable(name, address.getIp(), + new V1ServicePort().port(port.getPort()), payload); + if (d != null) { + discoverables.add(d); + } + } + } + } + } + } + } + } catch (Exception e) { + LOG.warn("Failed to retrieve endpoints for service {}, falling back to service hostname", name, e); + } + + if (!discoverables.isEmpty()) { + return discoverables; } - // We don't expect there is more than one service port, hence only pick the first one + String hostname = String.format("%s.%s", meta.getName(), namespace); return servicePorts.stream() .map(port -> createDiscoverable( name, hostname, diff --git a/task-manager-service.yaml b/task-manager-service.yaml index d14a24b60d35..fc5491ff767e 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -51,6 +51,8 @@ spec: - mountPath: /etc/cdap/security name: cdap-security readOnly: true + - mountPath: /cdap_configmap + name: cdap-cm-vol-cdap-sidhdirenge-jun8-configmap volumes: - downwardAPI: defaultMode: 420 @@ -80,6 +82,10 @@ spec: secret: defaultMode: 420 secretName: cdap-security + - configMap: + defaultMode: 420 + name: cdap-sidhdirenge-jun8-configmap + name: cdap-cm-vol-cdap-sidhdirenge-jun8-configmap --- apiVersion: v1 kind: Service From 912ce60d3144695ae501abc997d90e1f0a9ccefc Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 17 Jul 2026 11:52:59 +0000 Subject: [PATCH 06/54] feat: Implement Netty Proxy streaming skeleton for Task Manager POC - Replaced TaskManagerHttpHandler with standalone Netty ServerBootstrap - Added ProxyFrontendHandler for zero-copy (.retain()) chunk streaming and connection queuing - Added ProxyBackendHandler for worker response relay and reverse backpressure --- .../internal/remote/ProxyBackendHandler.java | 64 ++++++++ .../internal/remote/ProxyFrontendHandler.java | 152 ++++++++++++++++++ .../internal/remote/TaskManagerService.java | 81 +++++++--- .../remote/TaskManagerServiceModule.java | 5 +- 4 files changed, 275 insertions(+), 27 deletions(-) create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java new file mode 100644 index 000000000000..ad4b0488522c --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -0,0 +1,64 @@ +/* + * 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; + +public class ProxyBackendHandler extends ChannelInboundHandlerAdapter { + + private final Channel inboundChannel; + + public ProxyBackendHandler(Channel inboundChannel) { + this.inboundChannel = inboundChannel; + } + + @Override + public void channelRead(ChannelHandlerContext ctx, Object msg) { + // 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()) { + boolean isWritable = ctx.channel().isWritable(); + inboundChannel.config().setAutoRead(isWritable); + } + ctx.fireChannelWritabilityChanged(); + } + + @Override + public void channelInactive(ChannelHandlerContext ctx) { + ProxyFrontendHandler.closeOnFlush(inboundChannel); + } + + @Override + public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { + cause.printStackTrace(); + ProxyFrontendHandler.closeOnFlush(ctx.channel()); + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java new file mode 100644 index 000000000000..b4e8dd12c775 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -0,0 +1,152 @@ +/* + * 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.bootstrap.Bootstrap; +import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelFutureListener; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelInboundHandlerAdapter; +import io.netty.channel.ChannelInitializer; +import io.netty.channel.ChannelOption; +import io.netty.channel.ChannelPipeline; +import io.netty.channel.socket.SocketChannel; +import io.netty.channel.socket.nio.NioSocketChannel; +import io.netty.handler.codec.http.HttpClientCodec; +import io.netty.handler.codec.http.HttpContent; +import io.netty.handler.codec.http.HttpRequest; +import io.netty.util.ReferenceCountUtil; + +import java.util.LinkedList; +import java.util.Map; +import java.util.Queue; + +public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { + + private final Map workerPartitions; + private Channel outboundChannel; + private boolean connecting = false; + private final Queue pendingMessages = new LinkedList<>(); + + public ProxyFrontendHandler(Map workerPartitions) { + this.workerPartitions = workerPartitions; + } + + @Override + public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { + if (msg instanceof HttpRequest) { + HttpRequest req = (HttpRequest) msg; + + String partitionHeader = req.headers().get("X-Partition-ID"); + int partitionId = (partitionHeader != null) ? Integer.parseInt(partitionHeader) : 1; + String workerAddress = workerPartitions.getOrDefault(partitionId, "127.0.0.1:8081"); + String[] hostPort = workerAddress.split(":"); + + // Apply backpressure on client until connection established + ctx.channel().config().setAutoRead(false); + connecting = true; + + Bootstrap b = new Bootstrap(); + b.group(ctx.channel().eventLoop()) + .channel(NioSocketChannel.class) + .option(ChannelOption.SO_KEEPALIVE, true) + .handler(new ChannelInitializer() { + @Override + protected void initChannel(SocketChannel ch) { + ChannelPipeline p = ch.pipeline(); + p.addLast(new HttpClientCodec()); + p.addLast(new ProxyBackendHandler(ctx.channel())); + } + }); + + ChannelFuture f = b.connect(hostPort[0], Integer.parseInt(hostPort[1])); + outboundChannel = f.channel(); + + f.addListener((ChannelFutureListener) future -> { + connecting = false; + if (future.isSuccess()) { + // Drain pending messages queue + Object pendingMsg = pendingMessages.poll(); + while (pendingMsg != null) { + outboundChannel.write(pendingMsg); + pendingMsg = pendingMessages.poll(); + } + outboundChannel.flush(); + // Resume reading from client once connected + ctx.channel().config().setAutoRead(true); + } else { + // Release all pending messages + Object pendingMsg = pendingMessages.poll(); + while (pendingMsg != null) { + ReferenceCountUtil.release(pendingMsg); + pendingMsg = pendingMessages.poll(); + } + ctx.channel().close(); + } + }); + + pendingMessages.add(ReferenceCountUtil.retain(msg)); + + } else if (msg instanceof HttpContent) { + if (connecting) { + // Queue chunks while connection is establishing + pendingMessages.add(ReferenceCountUtil.retain(msg)); + } else if (outboundChannel != null && outboundChannel.isActive()) { + outboundChannel.writeAndFlush(ReferenceCountUtil.retain(msg)); + } + } + } + + @Override + public void channelReadComplete(ChannelHandlerContext ctx) { + if (outboundChannel != null && outboundChannel.isActive() && !connecting) { + outboundChannel.flush(); + } + ctx.fireChannelReadComplete(); + } + + @Override + public void channelWritabilityChanged(ChannelHandlerContext ctx) { + // App Fabric client channel is saturated; pause reading from backend Worker + if (outboundChannel != null && outboundChannel.isActive()) { + boolean isWritable = ctx.channel().isWritable(); + outboundChannel.config().setAutoRead(isWritable); + } + ctx.fireChannelWritabilityChanged(); + } + + @Override + public void channelInactive(ChannelHandlerContext ctx) { + if (outboundChannel != null) { + closeOnFlush(outboundChannel); + } + } + + @Override + public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { + cause.printStackTrace(); + closeOnFlush(ctx.channel()); + } + + static void closeOnFlush(Channel ch) { + if (ch.isActive()) { + ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE); + } + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 37510d9d7d3f..07f46a2b3749 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -19,49 +19,84 @@ import com.google.common.util.concurrent.AbstractIdleService; import com.google.inject.Inject; import io.cdap.cdap.common.conf.CConfiguration; -import io.cdap.cdap.common.http.CommonNettyHttpServiceFactory; -import io.cdap.http.NettyHttpService; +import io.netty.bootstrap.ServerBootstrap; +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelInitializer; +import io.netty.channel.ChannelPipeline; +import io.netty.channel.EventLoopGroup; +import io.netty.channel.nio.NioEventLoopGroup; +import io.netty.channel.socket.SocketChannel; +import io.netty.channel.socket.nio.NioServerSocketChannel; +import io.netty.handler.codec.http.HttpServerCodec; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Collections; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; /** - * Guice-managed service that runs the Centralized Task Manager HTTP Server. + * Guice-managed service that runs the Centralized Task Manager HTTP Server (Netty Proxy POC). */ public class TaskManagerService extends AbstractIdleService { private static final Logger LOG = LoggerFactory.getLogger(TaskManagerService.class); - private final NettyHttpService httpService; + + private final int port; + private final String address; + private EventLoopGroup bossGroup; + private EventLoopGroup workerGroup; + private ChannelFuture channelFuture; + + private final Map workerPartitions = new ConcurrentHashMap<>(); @Inject - TaskManagerService(CConfiguration cConf, - CommonNettyHttpServiceFactory commonNettyHttpServiceFactory, - TaskManagerHttpHandler taskManagerHttpHandler) { - - int port = cConf.getInt("task.manager.port", 11025); - String address = cConf.get("task.manager.address", "0.0.0.0"); + TaskManagerService(CConfiguration cConf) { + this.port = cConf.getInt("task.manager.port", 11025); + this.address = cConf.get("task.manager.address", "0.0.0.0"); - LOG.info("sidhdirenge - Initializing TaskManagerService on {}:{}", address, port); + // Mocking worker partitions for POC + workerPartitions.put(1, "127.0.0.1:8081"); + workerPartitions.put(2, "127.0.0.1:8082"); - this.httpService = commonNettyHttpServiceFactory.builder("task-manager", false) - .setHost(address) - .setPort(port) - .setHttpHandlers(Collections.singletonList(taskManagerHttpHandler)) - .build(); + LOG.info("sidhdirenge - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); } @Override protected void startUp() throws Exception { - LOG.info("sidhdirenge - Starting TaskManagerService HTTP server..."); - httpService.start(); - LOG.info("sidhdirenge - TaskManagerService HTTP server started successfully at {}", httpService.getBindAddress()); + LOG.info("sidhdirenge - Starting TaskManagerService Proxy HTTP server..."); + + bossGroup = new NioEventLoopGroup(1); + workerGroup = new NioEventLoopGroup(); + + ServerBootstrap b = new ServerBootstrap(); + b.group(bossGroup, workerGroup) + .channel(NioServerSocketChannel.class) + .childHandler(new ChannelInitializer() { + @Override + protected void initChannel(SocketChannel ch) { + ChannelPipeline p = ch.pipeline(); + p.addLast(new HttpServerCodec()); + // NOTICE: NO HttpObjectAggregator here! + p.addLast(new ProxyFrontendHandler(workerPartitions)); + } + }); + + channelFuture = b.bind(address, port).sync(); + LOG.info("sidhdirenge - TaskManagerService Proxy HTTP server started successfully at {}:{}", address, port); } @Override protected void shutDown() throws Exception { - LOG.info("sidhdirenge - Stopping TaskManagerService HTTP server..."); - httpService.stop(); - LOG.info("sidhdirenge - TaskManagerService HTTP server stopped."); + LOG.info("sidhdirenge - Stopping TaskManagerService Proxy HTTP server..."); + if (channelFuture != null) { + channelFuture.channel().close().sync(); + } + if (bossGroup != null) { + bossGroup.shutdownGracefully(); + } + if (workerGroup != null) { + workerGroup.shutdownGracefully(); + } + LOG.info("sidhdirenge - TaskManagerService Proxy HTTP server stopped."); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java index fb470462ab7e..c7cb33271a6e 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java @@ -29,10 +29,7 @@ protected void configure() { // Bind the core TaskManager as a singleton bind(TaskManager.class).toProvider(TaskManager::getInstance).in(Scopes.SINGLETON); - // Bind the HTTP handler - bind(TaskManagerHttpHandler.class).in(Scopes.SINGLETON); - - // Bind the service itself + // Bind the Netty Proxy service itself bind(TaskManagerService.class).in(Scopes.SINGLETON); } } From 3d94f45bae61e0165daa55c7ef3138b53850c445 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Mon, 20 Jul 2026 10:07:40 +0000 Subject: [PATCH 07/54] feat(poc): Implement Netty Proxy and Gatekeeper IAM reaper - Update Proxy to use thread-safe single PodState map with synchronized blocks - Add StickyLeaseManager lifecycle hooks for IAM Context injection - Remove per-task IAM sidecar wipes to enable Warm Boot leasing --- .../cdap/common/internal/remote/PodState.java | 46 +++++++++++ .../internal/remote/ProxyBackendHandler.java | 38 +++++++++- .../internal/remote/ProxyFrontendHandler.java | 76 +++++++++++++++---- .../internal/remote/StickyLeaseManager.java | 17 ++++- .../internal/remote/TaskManagerService.java | 8 +- .../remote/TaskWorkerHttpHandlerInternal.java | 28 ++++--- 6 files changed, 179 insertions(+), 34 deletions(-) create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java new file mode 100644 index 000000000000..6275e09f07e4 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -0,0 +1,46 @@ +/* + * 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; + + public PodState(String leasedNamespace, int inflightRequests) { + this.leasedNamespace = leasedNamespace; + this.inflightRequests = inflightRequests; + } + + 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; + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index ad4b0488522c..6672decdd487 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -20,17 +20,50 @@ 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; public class ProxyBackendHandler extends ChannelInboundHandlerAdapter { private final Channel inboundChannel; + private final Map podRegistry; + private final String targetWorkerAddress; - public ProxyBackendHandler(Channel inboundChannel) { + public ProxyBackendHandler(Channel inboundChannel, Map podRegistry, String targetWorkerAddress) { this.inboundChannel = inboundChannel; + this.podRegistry = podRegistry; + this.targetWorkerAddress = targetWorkerAddress; } @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"); + + if (activeTasksStr != null) { + try { + state.setInflightRequests(Integer.parseInt(activeTasksStr)); + } catch (NumberFormatException e) { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); + } + } else { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); + } + + if (leasedNamespace != null) { + state.setLeasedNamespace(leasedNamespace); + } + } + } + } + // Forward worker responses directly back to the client inboundChannel.writeAndFlush(msg).addListener((ChannelFutureListener) future -> { if (future.isSuccess()) { @@ -45,8 +78,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { public void channelWritabilityChanged(ChannelHandlerContext ctx) { // Backend Worker channel is saturated; pause reading from App Fabric client if (inboundChannel != null && inboundChannel.isActive()) { - boolean isWritable = ctx.channel().isWritable(); - inboundChannel.config().setAutoRead(isWritable); + inboundChannel.config().setAutoRead(ctx.channel().isWritable()); } ctx.fireChannelWritabilityChanged(); } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index b4e8dd12c775..678c30506389 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -28,9 +28,13 @@ import io.netty.channel.ChannelPipeline; import io.netty.channel.socket.SocketChannel; import io.netty.channel.socket.nio.NioSocketChannel; +import io.netty.handler.codec.http.DefaultFullHttpResponse; +import io.netty.handler.codec.http.FullHttpResponse; import io.netty.handler.codec.http.HttpClientCodec; import io.netty.handler.codec.http.HttpContent; import io.netty.handler.codec.http.HttpRequest; +import io.netty.handler.codec.http.HttpResponseStatus; +import io.netty.handler.codec.http.HttpVersion; import io.netty.util.ReferenceCountUtil; import java.util.LinkedList; @@ -39,13 +43,13 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { - private final Map workerPartitions; + private final Map podRegistry; private Channel outboundChannel; private boolean connecting = false; private final Queue pendingMessages = new LinkedList<>(); - public ProxyFrontendHandler(Map workerPartitions) { - this.workerPartitions = workerPartitions; + public ProxyFrontendHandler(Map podRegistry) { + this.podRegistry = podRegistry; } @Override @@ -53,10 +57,49 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (msg instanceof HttpRequest) { HttpRequest req = (HttpRequest) msg; - String partitionHeader = req.headers().get("X-Partition-ID"); - int partitionId = (partitionHeader != null) ? Integer.parseInt(partitionHeader) : 1; - String workerAddress = workerPartitions.getOrDefault(partitionId, "127.0.0.1:8081"); - String[] hostPort = workerAddress.split(":"); + String targetNamespace = req.headers().get("X-CDF-Namespace"); + if (targetNamespace == null) targetNamespace = "default"; + + String targetWorkerAddress = null; + + // 1. Warm Match: Thread-safe scan specifically locking evaluation + for (Map.Entry entry : podRegistry.entrySet()) { + PodState state = entry.getValue(); + synchronized (state) { + if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { + targetWorkerAddress = entry.getKey(); + state.setInflightRequests(state.getInflightRequests() + 1); + break; + } + } + } + + // 2. Idle Choice: Thread-safe claim of an idle pod + if (targetWorkerAddress == null) { + for (Map.Entry entry : podRegistry.entrySet()) { + PodState state = entry.getValue(); + synchronized (state) { + if (state.getInflightRequests() == 0) { + targetWorkerAddress = entry.getKey(); + state.setLeasedNamespace(targetNamespace); + state.setInflightRequests(1); + break; + } + } + } + } + + // 3. Busy Rejection: All pods saturated + if (targetWorkerAddress == null) { + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + ReferenceCountUtil.release(msg); + return; + } + + final String chosenWorker = targetWorkerAddress; + String[] hostPort = targetWorkerAddress.split(":"); // Apply backpressure on client until connection established ctx.channel().config().setAutoRead(false); @@ -71,7 +114,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception protected void initChannel(SocketChannel ch) { ChannelPipeline p = ch.pipeline(); p.addLast(new HttpClientCodec()); - p.addLast(new ProxyBackendHandler(ctx.channel())); + p.addLast(new ProxyBackendHandler(ctx.channel(), podRegistry, chosenWorker)); } }); @@ -81,22 +124,26 @@ protected void initChannel(SocketChannel ch) { f.addListener((ChannelFutureListener) future -> { connecting = false; if (future.isSuccess()) { - // Drain pending messages queue Object pendingMsg = pendingMessages.poll(); while (pendingMsg != null) { outboundChannel.write(pendingMsg); pendingMsg = pendingMessages.poll(); } outboundChannel.flush(); - // Resume reading from client once connected ctx.channel().config().setAutoRead(true); } else { - // Release all pending messages Object pendingMsg = pendingMessages.poll(); while (pendingMsg != null) { ReferenceCountUtil.release(pendingMsg); pendingMsg = pendingMessages.poll(); } + // Thread-safe decrement on fallback + PodState fallbackState = podRegistry.get(chosenWorker); + if (fallbackState != null) { + synchronized (fallbackState) { + fallbackState.setInflightRequests(Math.max(0, fallbackState.getInflightRequests() - 1)); + } + } ctx.channel().close(); } }); @@ -105,10 +152,11 @@ protected void initChannel(SocketChannel ch) { } else if (msg instanceof HttpContent) { if (connecting) { - // Queue chunks while connection is establishing pendingMessages.add(ReferenceCountUtil.retain(msg)); } else if (outboundChannel != null && outboundChannel.isActive()) { outboundChannel.writeAndFlush(ReferenceCountUtil.retain(msg)); + } else { + ReferenceCountUtil.release(msg); } } } @@ -123,10 +171,8 @@ public void channelReadComplete(ChannelHandlerContext ctx) { @Override public void channelWritabilityChanged(ChannelHandlerContext ctx) { - // App Fabric client channel is saturated; pause reading from backend Worker if (outboundChannel != null && outboundChannel.isActive()) { - boolean isWritable = ctx.channel().isWritable(); - outboundChannel.config().setAutoRead(isWritable); + outboundChannel.config().setAutoRead(ctx.channel().isWritable()); } ctx.fireChannelWritabilityChanged(); } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index 12754c57c642..e56e8e5d7a43 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -20,6 +20,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import javax.annotation.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,15 +45,21 @@ public class StickyLeaseManager { private final AtomicInteger activeTaskCount = new AtomicInteger(0); private final AtomicInteger totalTasksProcessedInLease = new AtomicInteger(0); private volatile long lastActivityTimeMillis; + private final Consumer onLeaseAcquired; + private final Runnable onLeaseReleased; public StickyLeaseManager() { - this(10, 10); + this(10, 10, null, null); } - public StickyLeaseManager(int maxConcurrentTasks, int maxTasksPerLease) { + public StickyLeaseManager(int maxConcurrentTasks, int maxTasksPerLease, + Consumer onLeaseAcquired, + Runnable onLeaseReleased) { this.maxConcurrentTasks = maxConcurrentTasks; this.maxTasksPerLease = maxTasksPerLease; this.lastActivityTimeMillis = System.currentTimeMillis(); + this.onLeaseAcquired = onLeaseAcquired; + this.onLeaseReleased = onLeaseReleased; } /** @@ -73,6 +80,9 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant LOG.info( "Lease claimed by namespace '{}' (Tier: {}) in {}ms (Boot penalty entirely avoided)", namespace.getNamespace(), tier, elapsed); + if (onLeaseAcquired != null) { + onLeaseAcquired.accept(namespace); + } return AcquisitionStatus.SUCCESS; } } @@ -155,6 +165,9 @@ public synchronized void releaseLease(String reason) { totalTasksProcessedInLease.set(0); LOG.info("Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", oldNamespace.getNamespace(), reason); + if (onLeaseReleased != null) { + onLeaseReleased.run(); + } } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 07f46a2b3749..a954bf0e8109 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -47,7 +47,7 @@ public class TaskManagerService extends AbstractIdleService { private EventLoopGroup workerGroup; private ChannelFuture channelFuture; - private final Map workerPartitions = new ConcurrentHashMap<>(); + private final Map podRegistry = new ConcurrentHashMap<>(); @Inject TaskManagerService(CConfiguration cConf) { @@ -55,8 +55,8 @@ public class TaskManagerService extends AbstractIdleService { this.address = cConf.get("task.manager.address", "0.0.0.0"); // Mocking worker partitions for POC - workerPartitions.put(1, "127.0.0.1:8081"); - workerPartitions.put(2, "127.0.0.1:8082"); + podRegistry.put("127.0.0.1:8081", new PodState(null, 0)); + podRegistry.put("127.0.0.1:8082", new PodState(null, 0)); LOG.info("sidhdirenge - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); } @@ -77,7 +77,7 @@ protected void initChannel(SocketChannel ch) { ChannelPipeline p = ch.pipeline(); p.addLast(new HttpServerCodec()); // NOTICE: NO HttpObjectAggregator here! - p.addLast(new ProxyFrontendHandler(workerPartitions)); + p.addLast(new ProxyFrontendHandler(podRegistry)); } }); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java index b7726c46af23..1473b23055e2 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java @@ -121,7 +121,24 @@ public TaskWorkerHttpHandlerInternal(CConfiguration cConf, TaskWorker.USER_CODE_ISOLATION_ENABLED); this.concurrentRequestLimit = cConf.getInt(TaskWorker.REQUEST_LIMIT); int maxTasksPerLease = cConf.getInt("task.worker.lease.max.tasks", 10); - this.stickyLeaseManager = new StickyLeaseManager(concurrentRequestLimit, maxTasksPerLease); + this.stickyLeaseManager = new StickyLeaseManager(concurrentRequestLimit, maxTasksPerLease, + (namespaceId) -> { + try { + GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, cConf); + } catch (Exception e) { + LOG.warn("Failed to set GCP metadata task context for namespace {}", namespaceId, e); + } + }, + () -> { + try { + // Wipe Sidecar IAM Tokens (Identity / Security boundary) + GcpMetadataTaskContextUtil.clearGcpMetadataTaskContext(cConf); + // TODO: Wipe the JVM artifactCache to enforce disk reclaim boundary here + } catch (Exception e) { + LOG.warn("Failed to clear GCP metadata task context", e); + } + } + ); ScheduledExecutorService leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( Threads.createDaemonThreadFactory("lease-reclamation")); @@ -274,8 +291,6 @@ public void run(FullHttpRequest request, HttpResponder responder) { runningRequestCount.incrementAndGet(); try { - // set the GcpMetadataTaskContext before running the task. - GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, cConf); RunnableTaskContext runnableTaskContext = new RunnableTaskContext(runnableTaskRequest); runnableTaskLauncher.launchRunnableTask(runnableTaskContext); @@ -307,13 +322,6 @@ public void run(FullHttpRequest request, HttpResponder responder) { // Potentially ran user code, hence terminate the runner. taskCompletionConsumer.accept(false, new TaskDetails(metricsCollectionService, startTime, true, runnableTaskRequest)); - } finally { - // clear the GcpMetadataTaskContext after the task is completed. - try { - GcpMetadataTaskContextUtil.clearGcpMetadataTaskContext(cConf); - } catch (Exception e) { - LOG.warn("Failed to clear GCP metadata task context", e); - } } } From 3de38096afd9bdfe3dd7451818d6c2fb80278b64 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 06:01:35 +0000 Subject: [PATCH 08/54] feat(poc): Wire up App Fabric client to Netty Proxy - Modify RemoteClient to exclusively route to the new L7 proxy natively - Inject X-CDF-Namespace headers into all worker outbound traffic - Remove manual task /finish ping as Netty proxy implicitly intercepts worker state - Implement 60-second glass-break direct-fallback if Proxy is completely unreachable --- .../common/internal/remote/RemoteClient.java | 128 +----------------- .../internal/remote/RemoteTaskExecutor.java | 9 +- 2 files changed, 13 insertions(+), 124 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index 9bea26792ea0..94507dcd5c1b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -269,33 +269,7 @@ public void executeStreamingRequest(HttpRequest request) } } private void notifyTaskManagerFinished(String namespace, Discoverable pod, boolean rejected) { - try { - URL url = new URL(TASK_MANAGER_URL + "/v3/taskmanager/finish"); - TaskManagerHttpHandler.FinishRequest finishRequest = new TaskManagerHttpHandler.FinishRequest(); - - java.lang.reflect.Field nsField = finishRequest.getClass().getDeclaredField("namespace"); - nsField.setAccessible(true); - nsField.set(finishRequest, namespace); - - TaskManagerHttpHandler.PodInfo podInfo = new TaskManagerHttpHandler.PodInfo( - pod.getSocketAddress().getHostString(), pod.getSocketAddress().getPort()); - java.lang.reflect.Field podField = finishRequest.getClass().getDeclaredField("pod"); - podField.setAccessible(true); - podField.set(finishRequest, podInfo); - - java.lang.reflect.Field rejectedField = finishRequest.getClass().getDeclaredField("rejected"); - rejectedField.setAccessible(true); - rejectedField.set(finishRequest, rejected); - - HttpRequest req = HttpRequest.post(url) - .addHeader(HttpHeaders.CONTENT_TYPE, "application/json") - .withBody(GSON.toJson(finishRequest)) - .build(); - - HttpRequests.execute(req, httpRequestConfig); - } catch (Exception e) { - LOG.warn("sidhdirenge - Failed to notify Task Manager of task completion", e); - } + // POC: No longer needed! The Netty Proxy syncs state implicitly via HTTP response interception. } /** @@ -368,104 +342,12 @@ public URL resolve(String resource, @Nullable String routingKey) { } } - LOG.info("sidhdirenge - RemoteClient resolving stickily via TaskManager for service {} with routingKey: {}", - discoverableServiceName, routingKey); - - // 1. Fetch all currently discovered endpoints - Iterable discoverables = () -> discoveryClient.discover(discoverableServiceName) - .iterator(); - List list = new ArrayList<>(); - for (Discoverable d : discoverables) { - // Perform DNS lookup to resolve the service hostname into individual pod IPs (for headless services) - try { - java.net.InetAddress[] addresses = java.net.InetAddress.getAllByName( - d.getSocketAddress().getHostName()); - for (java.net.InetAddress addr : addresses) { - list.add(new Discoverable(d.getName(), - new java.net.InetSocketAddress(addr.getHostAddress(), d.getSocketAddress().getPort()), - d.getPayload())); - } - } catch (java.net.UnknownHostException e) { - // Fallback to original discoverable if DNS lookup fails - list.add(d); - } - } - - if (list.isEmpty()) { - throw new ServiceUnavailableException(discoverableServiceName); - } - - // 2. Sort endpoints by IP address and port to ensure consistent ordering across all client instances - list.sort(Comparator.comparing((Discoverable d) -> d.getSocketAddress().getHostName()) - .thenComparingInt(d -> d.getSocketAddress().getPort())); - - // 3. Delegate to the standalone TaskManager Service over HTTP - Discoverable discoverable = null; - try { - URL url = new URL(TASK_MANAGER_URL + "/v3/taskmanager/resolve"); - TaskManagerHttpHandler.ResolveRequest resolveRequest = new TaskManagerHttpHandler.ResolveRequest(); - - java.lang.reflect.Field nsField = resolveRequest.getClass().getDeclaredField("namespace"); - nsField.setAccessible(true); - nsField.set(resolveRequest, routingKey); - - List podInfos = new ArrayList<>(); - for (Discoverable pod : list) { - podInfos.add(new TaskManagerHttpHandler.PodInfo( - pod.getSocketAddress().getHostString(), pod.getSocketAddress().getPort())); - } - java.lang.reflect.Field podsField = resolveRequest.getClass().getDeclaredField("pods"); - podsField.setAccessible(true); - podsField.set(resolveRequest, podInfos); - - HttpRequest req = HttpRequest.post(url) - .addHeader(HttpHeaders.CONTENT_TYPE, "application/json") - .withBody(GSON.toJson(resolveRequest)) - .build(); - - HttpResponse resp = HttpRequests.execute(req, httpRequestConfig); - if (resp.getResponseCode() == HttpURLConnection.HTTP_OK) { - TaskManagerHttpHandler.PodInfo selectedPodInfo = GSON.fromJson( - resp.getResponseBodyAsString(), TaskManagerHttpHandler.PodInfo.class); - - byte[] payload = list.isEmpty() ? new byte[0] : list.get(0).getPayload(); - for (Discoverable d : list) { - if (d.getSocketAddress().getPort() == selectedPodInfo.getPort() - && (d.getSocketAddress().getHostName().equals(selectedPodInfo.getHost()) - || (d.getSocketAddress().getAddress() != null - && d.getSocketAddress().getAddress().getHostAddress().equals(selectedPodInfo.getHost())))) { - payload = d.getPayload(); - break; - } - } - discoverable = new Discoverable("task.worker", - new java.net.InetSocketAddress(selectedPodInfo.getHost(), selectedPodInfo.getPort()), payload); - } - } catch (Exception e) { - LOG.warn("sidhdirenge - Failed to resolve pod via Task Manager HTTP Service. Falling back to local hashing.", e); - } - - // Fallback: If Task Manager is down or returns error, use standard consistent hashing - if (discoverable == null) { - int baseIndex = (routingKey.hashCode() & Integer.MAX_VALUE) % list.size(); - discoverable = list.get(baseIndex); - LOG.warn("sidhdirenge - TaskManager resolution failed. Falling back to default index {}", baseIndex); - } - - // Store resolved pod context in ThreadLocal for task execution callbacks - CURRENT_RESOLVED_POD.set(discoverable); - CURRENT_ROUTING_KEY.set(routingKey); - - LOG.info("sidhdirenge - Centralized TaskManager selected warm pod IP {} for routingKey: {}", - discoverable.getSocketAddress(), routingKey); - - URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); + LOG.info("sidhdirenge - RemoteClient routing directly to Netty TaskManager L7 proxy for routingKey: {}", routingKey); try { - return rewriteUrl(uri.toURL()); + String cleanPath = (basePath + resource).replaceAll("//+", "/"); + return new URL(TASK_MANAGER_URL + "/" + cleanPath); } catch (MalformedURLException e) { - throw new IllegalStateException( - String.format("Discovered service %s, but it announced malformed URL %s", - discoverableServiceName, uri), e); + throw new ServiceUnavailableException(discoverableServiceName, e); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index aa4682c8fec2..20dd1b27a593 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -143,8 +143,15 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception namespace); } } + String routingKey = namespace; + if (System.currentTimeMillis() - startTime > 60000) { + LOG.warn("sidhdirenge - TaskManager Proxy unreachable for 60s! Bypassing proxy and falling back to direct Worker routing!"); + routingKey = null; // Setting to null triggers CDAP's native RandomEndpoint discovery in RemoteClient + } + HttpRequest.Builder requestBuilder = remoteClient - .requestBuilder(HttpMethod.POST, workerUrl, namespace) + .requestBuilder(HttpMethod.POST, workerUrl, routingKey) + .addHeader("X-CDF-Namespace", namespace) .withBody(requestBody.duplicate()); if (compression) { requestBuilder.addHeader(HttpHeaders.CONTENT_ENCODING, "gzip"); From 92430a88483348aa0805e0bcc65855d653f6024e Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 06:40:29 +0000 Subject: [PATCH 09/54] feat(poc): Wire up Netty Proxy to K8s Discovery service - Delete background polling thread in TaskManagerService - Inject DiscoveryServiceClient directly into ProxyFrontendHandler - Execute synchronous, real-time ZK watch evaluations directly on the EventLoop using Twill's local cache - Resolve actual active Kubernetes Pod IPs instead of mocked local ports --- .../internal/remote/ProxyFrontendHandler.java | 22 ++++++++++++++++++- .../internal/remote/TaskManagerService.java | 22 ++++++++++++------- 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 678c30506389..ac035618b664 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -40,16 +40,23 @@ import java.util.LinkedList; import java.util.Map; import java.util.Queue; +import java.util.Set; +import java.util.HashSet; +import org.apache.twill.discovery.Discoverable; +import org.apache.twill.discovery.DiscoveryServiceClient; +import io.cdap.cdap.common.conf.Constants; public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { private final Map podRegistry; + private final DiscoveryServiceClient discoveryServiceClient; private Channel outboundChannel; private boolean connecting = false; private final Queue pendingMessages = new LinkedList<>(); - public ProxyFrontendHandler(Map podRegistry) { + public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { this.podRegistry = podRegistry; + this.discoveryServiceClient = discoveryServiceClient; } @Override @@ -57,6 +64,19 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (msg instanceof HttpRequest) { HttpRequest req = (HttpRequest) msg; + // 0. Synchronous K8s Discovery (Zero-Stale State) + // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient evaluates a local memory cache backed by a push-based ZooKeeper watch. + Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); + Set activePods = new HashSet<>(); + for (Discoverable d : discoverables) { + activePods.add(d.getSocketAddress().getHostString() + ":" + d.getSocketAddress().getPort()); + } + + for (String podIp : activePods) { + podRegistry.putIfAbsent(podIp, new PodState(null, 0)); + } + podRegistry.keySet().removeIf(existingPod -> !activePods.contains(existingPod)); + String targetNamespace = req.headers().get("X-CDF-Namespace"); if (targetNamespace == null) targetNamespace = "default"; diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index a954bf0e8109..2e50677c60b1 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -31,6 +31,11 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.apache.twill.discovery.Discoverable; +import org.apache.twill.discovery.DiscoveryServiceClient; +import io.cdap.cdap.common.conf.Constants; +import java.util.Set; +import java.util.HashSet; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -49,14 +54,13 @@ public class TaskManagerService extends AbstractIdleService { private final Map podRegistry = new ConcurrentHashMap<>(); + private final DiscoveryServiceClient discoveryServiceClient; + @Inject - TaskManagerService(CConfiguration cConf) { + TaskManagerService(CConfiguration cConf, DiscoveryServiceClient discoveryServiceClient) { this.port = cConf.getInt("task.manager.port", 11025); this.address = cConf.get("task.manager.address", "0.0.0.0"); - - // Mocking worker partitions for POC - podRegistry.put("127.0.0.1:8081", new PodState(null, 0)); - podRegistry.put("127.0.0.1:8082", new PodState(null, 0)); + this.discoveryServiceClient = discoveryServiceClient; LOG.info("sidhdirenge - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); } @@ -65,8 +69,10 @@ public class TaskManagerService extends AbstractIdleService { protected void startUp() throws Exception { LOG.info("sidhdirenge - Starting TaskManagerService Proxy HTTP server..."); - bossGroup = new NioEventLoopGroup(1); - workerGroup = new NioEventLoopGroup(); + bossGroup = new NioEventLoopGroup(1, + new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-boss-thread-%d").build()); + workerGroup = new NioEventLoopGroup(0, + new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-worker-thread-%d").build()); ServerBootstrap b = new ServerBootstrap(); b.group(bossGroup, workerGroup) @@ -77,7 +83,7 @@ protected void initChannel(SocketChannel ch) { ChannelPipeline p = ch.pipeline(); p.addLast(new HttpServerCodec()); // NOTICE: NO HttpObjectAggregator here! - p.addLast(new ProxyFrontendHandler(podRegistry)); + p.addLast(new ProxyFrontendHandler(podRegistry, discoveryServiceClient)); } }); From f4c54eb699ed6bfe467efdbc7bbcaa5ea4c575ff Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 07:39:29 +0000 Subject: [PATCH 10/54] feat(poc): Inject ground truth occupancy headers into Task Worker responses - Expose X-Active-Tasks and X-Leased-Namespace from StickyLeaseManager - Ensure headers populate on both successful launches and lease rejections (429/409) - Enables Netty Gatekeeper to perform self-healing routing map corrections autonomously --- .../remote/TaskWorkerHttpHandlerInternal.java | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java index 1473b23055e2..43326922c7b5 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java @@ -284,7 +284,15 @@ public void run(FullHttpRequest request, HttpResponder responder) { if (leaseStatus != StickyLeaseManager.AcquisitionStatus.SUCCESS) { LOG.warn("Rejecting request for namespace {} due to lease status: {}", namespaceId, leaseStatus); - responder.sendStatus(HttpResponseStatus.TOO_MANY_REQUESTS); + + String currentLease = stickyLeaseManager.getCurrentLease() != null ? + stickyLeaseManager.getCurrentLease().getNamespace() : ""; + + responder.sendString(HttpResponseStatus.TOO_MANY_REQUESTS, + "Rejected due to lease status: " + leaseStatus.name(), + new DefaultHttpHeaders() + .add("X-Active-Tasks", String.valueOf(stickyLeaseManager.getActiveTaskCount())) + .add("X-Leased-Namespace", currentLease)); return; } @@ -298,11 +306,16 @@ public void run(FullHttpRequest request, HttpResponder responder) { startTime, runnableTaskContext.isTerminateOnComplete(), runnableTaskRequest); + String currentLease = stickyLeaseManager.getCurrentLease() != null ? + stickyLeaseManager.getCurrentLease().getNamespace() : ""; + responder.sendContent(HttpResponseStatus.OK, new RunnableTaskBodyProducer(runnableTaskContext, taskCompletionConsumer, taskDetails), - new DefaultHttpHeaders().add(HttpHeaders.CONTENT_TYPE, - MediaType.APPLICATION_OCTET_STREAM)); + new DefaultHttpHeaders() + .add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_OCTET_STREAM) + .add("X-Active-Tasks", String.valueOf(stickyLeaseManager.getActiveTaskCount())) + .add("X-Leased-Namespace", currentLease)); } catch (ClassNotFoundException | ClassCastException ex) { responder.sendString(HttpResponseStatus.BAD_REQUEST, exceptionToJson(ex), From 45555ad347e69d3ee46dd97b9735f111158af6ee Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 09:34:15 +0000 Subject: [PATCH 11/54] chore: update logging prefix from sidhdirenge to shruzard --- .../common/internal/remote/RemoteClient.java | 2 +- .../internal/remote/RemoteTaskExecutor.java | 4 ++-- .../cdap/common/internal/remote/TaskManager.java | 10 +++++----- .../internal/remote/TaskManagerService.java | 10 +++++----- .../remote/TaskWorkerHttpHandlerInternal.java | 2 +- task-manager-service.yaml | 16 ++++++++-------- 6 files changed, 22 insertions(+), 22 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index 94507dcd5c1b..e79913bfa7da 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -342,7 +342,7 @@ public URL resolve(String resource, @Nullable String routingKey) { } } - LOG.info("sidhdirenge - RemoteClient routing directly to Netty TaskManager L7 proxy for routingKey: {}", routingKey); + LOG.info("shruzard - RemoteClient routing directly to Netty TaskManager L7 proxy for routingKey: {}", routingKey); try { String cleanPath = (basePath + resource).replaceAll("//+", "/"); return new URL(TASK_MANAGER_URL + "/" + cleanPath); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 20dd1b27a593..9847d8ffde64 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -139,13 +139,13 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception String embeddedNamespace = runnableTaskRequest.getParam().getEmbeddedTaskRequest().getNamespace(); if (embeddedNamespace != null && !embeddedNamespace.isEmpty()) { namespace = embeddedNamespace; - LOG.info("sidhdirenge - RemoteTaskExecutor: Mapped SystemAppTask namespace to embedded: {}", + LOG.info("shruzard - RemoteTaskExecutor: Mapped SystemAppTask namespace to embedded: {}", namespace); } } String routingKey = namespace; if (System.currentTimeMillis() - startTime > 60000) { - LOG.warn("sidhdirenge - TaskManager Proxy unreachable for 60s! Bypassing proxy and falling back to direct Worker routing!"); + LOG.warn("shruzard - TaskManager Proxy unreachable for 60s! Bypassing proxy and falling back to direct Worker routing!"); routingKey = null; // Setting to null triggers CDAP's native RandomEndpoint discovery in RemoteClient } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java index 0ab82fe9967f..811cb514b9f3 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java @@ -99,7 +99,7 @@ public Discoverable resolvePod(String namespace, List availablePod podActiveTaskCounts.put(podIp, 0); podTotalTaskProcessedCounts.put(podIp, 0); - LOG.info("sidhdirenge - TaskManager: Established new lease for namespace '{}' on pod '{}'", + LOG.info("shruzard - TaskManager: Established new lease for namespace '{}' on pod '{}'", namespace, podIp); leasedPodIp = podIp; @@ -111,7 +111,7 @@ public Discoverable resolvePod(String namespace, List availablePod // 3. Fallback: If all pods are leased to other namespaces, find the pod with the least load if (selectedPod == null) { - LOG.warn("sidhdirenge - TaskManager: All pods are leased. Falling back to least-loaded pod."); + LOG.warn("shruzard - TaskManager: All pods are leased. Falling back to least-loaded pod."); int minLoad = Integer.MAX_VALUE; List bestPods = new ArrayList<>(); for (Discoverable pod : availablePods) { @@ -146,7 +146,7 @@ public Discoverable resolvePod(String namespace, List availablePod podActiveTaskCounts.put(leasedPodIp, activeTasks); podTotalTaskProcessedCounts.put(leasedPodIp, totalProcessed); - LOG.info("sidhdirenge - TaskManager: Routing task for '{}' to pod '{}' (Active: {}, Total: {})", + LOG.info("shruzard - TaskManager: Routing task for '{}' to pod '{}' (Active: {}, Total: {})", namespace, leasedPodIp, activeTasks, totalProcessed); } @@ -174,13 +174,13 @@ public void finishTask(String namespace, Discoverable pod, boolean rejected) { } else { int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); if (totalProcessed >= MAX_TOTAL_TASKS_BEFORE_RESET && activeTasks == 0) { - LOG.info("sidhdirenge - TaskManager: Pod '{}' finished all active tasks " + LOG.info("shruzard - TaskManager: Pod '{}' finished all active tasks " + "after reaching reset threshold. Reclaiming lease.", podIp); releaseLease(podIp); } } - LOG.info("sidhdirenge - TaskManager: Task finished for '{}' on pod '{}' (Remaining active: {}, Rejected: {})", + LOG.info("shruzard - TaskManager: Task finished for '{}' on pod '{}' (Remaining active: {}, Rejected: {})", namespace, podIp, activeTasks, rejected); } finally { lock.unlock(); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 2e50677c60b1..346605e8d9ae 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -62,12 +62,12 @@ public class TaskManagerService extends AbstractIdleService { this.address = cConf.get("task.manager.address", "0.0.0.0"); this.discoveryServiceClient = discoveryServiceClient; - LOG.info("sidhdirenge - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); + LOG.info("shruzard - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); } @Override protected void startUp() throws Exception { - LOG.info("sidhdirenge - Starting TaskManagerService Proxy HTTP server..."); + LOG.info("shruzard - Starting TaskManagerService Proxy HTTP server..."); bossGroup = new NioEventLoopGroup(1, new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-boss-thread-%d").build()); @@ -88,12 +88,12 @@ protected void initChannel(SocketChannel ch) { }); channelFuture = b.bind(address, port).sync(); - LOG.info("sidhdirenge - TaskManagerService Proxy HTTP server started successfully at {}:{}", address, port); + LOG.info("shruzard - TaskManagerService Proxy HTTP server started successfully at {}:{}", address, port); } @Override protected void shutDown() throws Exception { - LOG.info("sidhdirenge - Stopping TaskManagerService Proxy HTTP server..."); + LOG.info("shruzard - Stopping TaskManagerService Proxy HTTP server..."); if (channelFuture != null) { channelFuture.channel().close().sync(); } @@ -103,6 +103,6 @@ protected void shutDown() throws Exception { if (workerGroup != null) { workerGroup.shutdownGracefully(); } - LOG.info("sidhdirenge - TaskManagerService Proxy HTTP server stopped."); + LOG.info("shruzard - TaskManagerService Proxy HTTP server stopped."); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java index 43326922c7b5..24666e4d3044 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java @@ -245,7 +245,7 @@ private void stopAndShutdown(ScheduledExecutorService executorService, Consumer< @POST @Path("/run") public void run(FullHttpRequest request, HttpResponder responder) { - LOG.info("sidhdirenge - Received task on worker {} for namespace :{}", + LOG.info("shruzard - Received task on worker {} for namespace :{}", System.getenv("HOSTNAME") != null ? System.getenv("HOSTNAME") : "unknown", request.headers()); if (mustRestart.get()) { diff --git a/task-manager-service.yaml b/task-manager-service.yaml index fc5491ff767e..00ab201bcd84 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -5,7 +5,7 @@ metadata: namespace: default labels: cdap.service: task.manager - cdap.instance: sidhdirenge-jun8 + cdap.instance: shruzard-jun8 spec: replicas: 1 selector: @@ -15,9 +15,9 @@ spec: metadata: labels: cdap.service: task.manager - cdap.instance: sidhdirenge-jun8 + cdap.instance: shruzard-jun8 spec: - serviceAccountName: cdap-sidhdirenge-jun8-system-sa + serviceAccountName: cdap-shruzard-jun8-system-sa containers: - name: task-manager image: us-east1-docker.pkg.dev/ld27be8c949817660-tp/ar-demo/cloud-data-fusion:latest @@ -52,7 +52,7 @@ spec: name: cdap-security readOnly: true - mountPath: /cdap_configmap - name: cdap-cm-vol-cdap-sidhdirenge-jun8-configmap + name: cdap-cm-vol-cdap-shruzard-jun8-configmap volumes: - downwardAPI: defaultMode: 420 @@ -72,11 +72,11 @@ spec: name: podinfo - configMap: defaultMode: 420 - name: cdap-sidhdirenge-jun8-cconf + name: cdap-shruzard-jun8-cconf name: cdap-conf - configMap: defaultMode: 420 - name: cdap-sidhdirenge-jun8-hconf + name: cdap-shruzard-jun8-hconf name: hadoop-conf - name: cdap-security secret: @@ -84,8 +84,8 @@ spec: secretName: cdap-security - configMap: defaultMode: 420 - name: cdap-sidhdirenge-jun8-configmap - name: cdap-cm-vol-cdap-sidhdirenge-jun8-configmap + name: cdap-shruzard-jun8-configmap + name: cdap-cm-vol-cdap-shruzard-jun8-configmap --- apiVersion: v1 kind: Service From 26a1d1bfc97f04529f4692d3d2aca24ae3dafdff Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 09:35:30 +0000 Subject: [PATCH 12/54] chore: add comprehensive logging across Netty POC components with shruzard prefix --- .../common/internal/remote/ProxyBackendHandler.java | 11 +++++++++++ .../common/internal/remote/ProxyFrontendHandler.java | 8 ++++++++ .../common/internal/remote/StickyLeaseManager.java | 6 +++--- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index 6672decdd487..6f9ca733805c 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -23,9 +23,14 @@ 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 podRegistry; private final String targetWorkerAddress; @@ -57,9 +62,15 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); } + if (leasedNamespace != null) { state.setLeasedNamespace(leasedNamespace); } + + if (activeTasksStr != null || leasedNamespace != null) { + LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState for {}. Occupancy: {}, Namespace: {}", + targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); + } } } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index ac035618b664..dc6301012d7d 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -45,9 +45,14 @@ import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; import io.cdap.cdap.common.conf.Constants; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { + private static final Logger LOG = LoggerFactory.getLogger(ProxyFrontendHandler.class); + + private final Map podRegistry; private final DiscoveryServiceClient discoveryServiceClient; private Channel outboundChannel; @@ -89,6 +94,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { targetWorkerAddress = entry.getKey(); state.setInflightRequests(state.getInflightRequests() + 1); + LOG.info("shruzard - ProxyFrontendHandler: Found warm match for '{}' at {}. Occupancy: {}", targetNamespace, targetWorkerAddress, state.getInflightRequests()); break; } } @@ -103,6 +109,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception targetWorkerAddress = entry.getKey(); state.setLeasedNamespace(targetNamespace); state.setInflightRequests(1); + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod at {} for namespace '{}'.", targetWorkerAddress, targetNamespace); break; } } @@ -111,6 +118,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // 3. Busy Rejection: All pods saturated if (targetWorkerAddress == null) { + LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased incorrectly. Rejecting request for namespace '{}'", targetNamespace); FullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index e56e8e5d7a43..65ffeb4c07e8 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -94,7 +94,7 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant } // Mismatching namespace -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) - LOG.info("Enforcement: Rejecting request for namespace '{}', current lease is held by '{}'", + LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}'", namespace.getNamespace(), currentLease.get()); return AcquisitionStatus.REJECTED_MISMATCH; } @@ -109,7 +109,7 @@ public synchronized AcquisitionStatus startTask(NamespaceId namespace, TenantTie } if (activeTaskCount.get() >= maxConcurrentTasks) { - LOG.info("Concurrency limit reached ({} tasks active) for namespace '{}'", + LOG.info("shruzard - StickyLeaseManager: Concurrency limit reached ({} tasks active) for namespace '{}'", activeTaskCount.get(), namespace.getNamespace()); return AcquisitionStatus.REJECTED_MAX_CONCURRENCY; } @@ -163,7 +163,7 @@ public synchronized void releaseLease(String reason) { if (oldNamespace != null) { activeTaskCount.set(0); totalTasksProcessedInLease.set(0); - LOG.info("Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", + LOG.info("shruzard - StickyLeaseManager: Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", oldNamespace.getNamespace(), reason); if (onLeaseReleased != null) { onLeaseReleased.run(); From 35983026092f057a766d09ab4815fa1b69bd7023 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 09:43:17 +0000 Subject: [PATCH 13/54] refactor(poc): eliminate legacy task confirmation network callbacks - Removed ThreadLocal pod tracking from App Fabric RemoteClient - Deleted notifyTaskManagerFinished endpoint and logic - Relies entirely on the Netty Proxy passive response header interception for state synchronization --- .../common/internal/remote/RemoteClient.java | 37 +------------------ 1 file changed, 1 insertion(+), 36 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index e79913bfa7da..3f93f327f5a0 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -67,8 +67,7 @@ public class RemoteClient { public static final String RUNTIME_SERVICE_ROUTING_BASE_URI = "cdap.runtime.service.routing.base.uri"; private static final Logger LOG = LoggerFactory.getLogger(RemoteClient.class); - private static final ThreadLocal CURRENT_RESOLVED_POD = new ThreadLocal<>(); - private static final ThreadLocal CURRENT_ROUTING_KEY = new ThreadLocal<>(); + private static final String TASK_MANAGER_URL = "http://cdap-task-manager.default.svc.cluster.local:11025"; private static final Gson GSON = new Gson(); @@ -181,13 +180,9 @@ private HttpResponse executeNonIdempotent(HttpRequest request) throws IOExceptio HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength()); - boolean rejected = false; try { HttpResponse response = HttpRequests.execute(httpRequest, httpRequestConfig); int responseCode = response.getResponseCode(); - if (responseCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - rejected = true; - } // 503 is always retryable. Other 5xx errors are retryable if the request is idempotent (handled in // RemoteClient#executeIdempotent(HttpRequest) if (responseCode == HttpURLConnection.HTTP_UNAVAILABLE) { @@ -215,19 +210,7 @@ private HttpResponse executeNonIdempotent(HttpRequest request) throws IOExceptio } return response; } catch (ConnectException e) { - rejected = true; throw new ServiceUnavailableException(discoverableServiceName, e); - } catch (IOException | RuntimeException e) { - rejected = true; - throw e; - } finally { - Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); - String routingKey = CURRENT_ROUTING_KEY.get(); - if (resolvedPod != null && routingKey != null) { - notifyTaskManagerFinished(routingKey, resolvedPod, rejected); - } - CURRENT_RESOLVED_POD.remove(); - CURRENT_ROUTING_KEY.remove(); } } @@ -242,35 +225,17 @@ public void executeStreamingRequest(HttpRequest request) HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength(), request.getConsumer()); - boolean rejected = false; try { HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { - if (httpResponse.getResponseCode() == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - rejected = true; - } throw new IOException( String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), httpResponse.getResponseCode())); } httpResponse.consumeContent(); - } catch (IOException | RuntimeException e) { - rejected = true; - throw e; - } finally { - Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); - String routingKey = CURRENT_ROUTING_KEY.get(); - if (resolvedPod != null && routingKey != null) { - notifyTaskManagerFinished(routingKey, resolvedPod, rejected); - } - CURRENT_RESOLVED_POD.remove(); - CURRENT_ROUTING_KEY.remove(); } } - private void notifyTaskManagerFinished(String namespace, Discoverable pod, boolean rejected) { - // POC: No longer needed! The Netty Proxy syncs state implicitly via HTTP response interception. - } /** * Opens a {@link HttpURLConnection} for the given resource path. From 6a4f79c254cee436921f1966c081684f9df3ac4a Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 10:24:34 +0000 Subject: [PATCH 14/54] refactor --- .../cdap/common/internal/remote/PodState.java | 10 + .../internal/remote/ProxyBackendHandler.java | 2 + .../internal/remote/ProxyFrontendHandler.java | 10 +- .../internal/remote/StickyLeaseManager.java | 2 +- .../common/internal/remote/TaskManager.java | 199 ------------------ .../remote/TaskManagerHttpHandler.java | 157 -------------- .../remote/TaskManagerServiceModule.java | 3 +- .../internal/remote/TaskManagerTest.java | 115 ---------- 8 files changed, 21 insertions(+), 477 deletions(-) delete mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java delete mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java delete mode 100644 cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java index 6275e09f07e4..20b3d1d2858b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -22,10 +22,12 @@ public class PodState { private String leasedNamespace; private int inflightRequests; + private long lastActivityTime; public PodState(String leasedNamespace, int inflightRequests) { this.leasedNamespace = leasedNamespace; this.inflightRequests = inflightRequests; + this.lastActivityTime = 0; // Instantly trigger predictions on boot } public String getLeasedNamespace() { @@ -43,4 +45,12 @@ public int getInflightRequests() { public void setInflightRequests(int inflightRequests) { this.inflightRequests = inflightRequests; } + + public long getLastActivityTime() { + return lastActivityTime; + } + + public void setLastActivityTime(long lastActivityTime) { + this.lastActivityTime = lastActivityTime; + } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index 6f9ca733805c..a9d6b77a7306 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -67,6 +67,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { state.setLeasedNamespace(leasedNamespace); } + state.setLastActivityTime(System.currentTimeMillis()); + if (activeTasksStr != null || leasedNamespace != null) { LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState for {}. Occupancy: {}, Namespace: {}", targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index dc6301012d7d..c00d7ddfa8d8 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -100,16 +100,20 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } } - // 2. Idle Choice: Thread-safe claim of an idle pod + // 2. Idle Choice: Thread-safe claim of an unleased pod, OR an expired pod (35s predicted timeout avoiding clock drift) if (targetWorkerAddress == null) { for (Map.Entry entry : podRegistry.entrySet()) { PodState state = entry.getValue(); synchronized (state) { - if (state.getInflightRequests() == 0) { + boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); + boolean isExpiredIdle = (state.getInflightRequests() == 0 && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); + + if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { targetWorkerAddress = entry.getKey(); state.setLeasedNamespace(targetNamespace); state.setInflightRequests(1); - LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod at {} for namespace '{}'.", targetWorkerAddress, targetNamespace); + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod (Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", + isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); break; } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index 65ffeb4c07e8..6002725b8ad3 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -78,7 +78,7 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant lastActivityTimeMillis = System.currentTimeMillis(); long elapsed = System.currentTimeMillis() - claimStartTime; LOG.info( - "Lease claimed by namespace '{}' (Tier: {}) in {}ms (Boot penalty entirely avoided)", + "shruzard Lease claimed by namespace '{}' (Tier: {}) in {}ms (Boot penalty entirely avoided)", namespace.getNamespace(), tier, elapsed); if (onLeaseAcquired != null) { onLeaseAcquired.accept(namespace); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java deleted file mode 100644 index 811cb514b9f3..000000000000 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java +++ /dev/null @@ -1,199 +0,0 @@ -/* - * 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 java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.locks.ReentrantLock; -import javax.annotation.Nullable; -import org.apache.twill.discovery.Discoverable; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Centralized Task Manager for orchestrating Warm Sticky Leases on Task Worker pods. - * This class coordinates leases, concurrency, and logical resets. - */ -public class TaskManager { - - private static final Logger LOG = LoggerFactory.getLogger(TaskManager.class); - private static final TaskManager INSTANCE = new TaskManager(); - - // Concurrency and task limits based on the design doc - private static final int MAX_CONCURRENT_TASKS_PER_POD = 10; - static final int MAX_TOTAL_TASKS_BEFORE_RESET = 10; - - private final ReentrantLock lock = new ReentrantLock(); - - // Lease state maps - // Pod IP/Key -> Namespace currently leased - private final Map podLeases = new HashMap<>(); - // Pod IP/Key -> Active concurrent task count - private final Map podActiveTaskCounts = new HashMap<>(); - // Pod IP/Key -> Total tasks processed on the current lease - private final Map podTotalTaskProcessedCounts = new HashMap<>(); - - public static TaskManager getInstance() { - return INSTANCE; - } - - private TaskManager() { - // Singleton - } - - /** - * Resolves the target warm pod for a given namespace based on the sticky lease model. - * - * @param namespace the namespace requesting execution - * @param availablePods the list of currently discovered pods - * @return the selected pod, or null if no pod is available - */ - @Nullable - public Discoverable resolvePod(String namespace, List availablePods) { - lock.lock(); - try { - String leasedPodIp = null; - Discoverable selectedPod = null; - - // 1. Find if a pod is already leased to this namespace and has capacity - for (Discoverable pod : availablePods) { - String podIp = getPodKey(pod); - String currentLease = podLeases.get(podIp); - - if (namespace.equals(currentLease)) { - int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); - if (activeTasks < MAX_CONCURRENT_TASKS_PER_POD && totalProcessed < MAX_TOTAL_TASKS_BEFORE_RESET) { - leasedPodIp = podIp; - selectedPod = pod; - break; - } - } - } - - // 2. If no active lease exists (or it is at capacity), find an idle/unleased pod - if (selectedPod == null) { - for (Discoverable pod : availablePods) { - String podIp = getPodKey(pod); - String currentLease = podLeases.get(podIp); - - if (currentLease == null) { - // Establish a new lease on this idle pod - podLeases.put(podIp, namespace); - podActiveTaskCounts.put(podIp, 0); - podTotalTaskProcessedCounts.put(podIp, 0); - - LOG.info("shruzard - TaskManager: Established new lease for namespace '{}' on pod '{}'", - namespace, podIp); - - leasedPodIp = podIp; - selectedPod = pod; - break; - } - } - } - - // 3. Fallback: If all pods are leased to other namespaces, find the pod with the least load - if (selectedPod == null) { - LOG.warn("shruzard - TaskManager: All pods are leased. Falling back to least-loaded pod."); - int minLoad = Integer.MAX_VALUE; - List bestPods = new ArrayList<>(); - for (Discoverable pod : availablePods) { - String podIp = getPodKey(pod); - int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); - if (activeTasks < minLoad) { - minLoad = activeTasks; - bestPods.clear(); - bestPods.add(pod); - } else if (activeTasks == minLoad) { - bestPods.add(pod); - } - } - - if (!bestPods.isEmpty()) { - int randomIndex = java.util.concurrent.ThreadLocalRandom.current().nextInt(bestPods.size()); - selectedPod = bestPods.get(randomIndex); - leasedPodIp = getPodKey(selectedPod); - - // Force-assign lease to the new namespace - podLeases.put(leasedPodIp, namespace); - podActiveTaskCounts.put(leasedPodIp, 0); - podTotalTaskProcessedCounts.put(leasedPodIp, 0); - } - } - - // 4. Increment task counts for the selected pod - if (selectedPod != null) { - int activeTasks = podActiveTaskCounts.getOrDefault(leasedPodIp, 0) + 1; - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(leasedPodIp, 0) + 1; - - podActiveTaskCounts.put(leasedPodIp, activeTasks); - podTotalTaskProcessedCounts.put(leasedPodIp, totalProcessed); - - LOG.info("shruzard - TaskManager: Routing task for '{}' to pod '{}' (Active: {}, Total: {})", - namespace, leasedPodIp, activeTasks, totalProcessed); - } - - return selectedPod; - } finally { - lock.unlock(); - } - } - - public void finishTask(String namespace, Discoverable pod, boolean rejected) { - lock.lock(); - try { - String podIp = getPodKey(pod); - int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); - if (activeTasks > 0) { - activeTasks = activeTasks - 1; - podActiveTaskCounts.put(podIp, activeTasks); - } - - if (rejected) { - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); - if (totalProcessed > 0) { - podTotalTaskProcessedCounts.put(podIp, totalProcessed - 1); - } - } else { - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); - if (totalProcessed >= MAX_TOTAL_TASKS_BEFORE_RESET && activeTasks == 0) { - LOG.info("shruzard - TaskManager: Pod '{}' finished all active tasks " - + "after reaching reset threshold. Reclaiming lease.", podIp); - releaseLease(podIp); - } - } - - LOG.info("shruzard - TaskManager: Task finished for '{}' on pod '{}' (Remaining active: {}, Rejected: {})", - namespace, podIp, activeTasks, rejected); - } finally { - lock.unlock(); - } - } - - private void releaseLease(String podIp) { - podLeases.remove(podIp); - podActiveTaskCounts.remove(podIp); - podTotalTaskProcessedCounts.remove(podIp); - } - - private String getPodKey(Discoverable pod) { - return pod.getSocketAddress().getHostString() + ":" + pod.getSocketAddress().getPort(); - } -} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java deleted file mode 100644 index d53ea147929b..000000000000 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java +++ /dev/null @@ -1,157 +0,0 @@ -/* - * 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 com.google.gson.Gson; -import com.google.gson.reflect.TypeToken; -import io.cdap.http.AbstractHttpHandler; -import io.cdap.http.HttpResponder; -import io.netty.handler.codec.http.FullHttpRequest; -import io.netty.handler.codec.http.HttpResponseStatus; -import org.apache.twill.discovery.Discoverable; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.lang.reflect.Type; -import java.net.InetSocketAddress; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.List; -import javax.ws.rs.POST; -import javax.ws.rs.Path; - -/** - * Netty HTTP Handler for the standalone Task Manager Service. - */ -@Path("/v3/taskmanager") -public class TaskManagerHttpHandler extends AbstractHttpHandler { - - private static final Logger LOG = LoggerFactory.getLogger(TaskManagerHttpHandler.class); - private static final Gson GSON = new Gson(); - private final TaskManager taskManager = TaskManager.getInstance(); - - @POST - @Path("/resolve") - public void resolve(FullHttpRequest request, HttpResponder responder) { - try { - String jsonBody = request.content().toString(StandardCharsets.UTF_8); - ResolveRequest resolveRequest = GSON.fromJson(jsonBody, ResolveRequest.class); - - if (resolveRequest == null || resolveRequest.getNamespace() == null || resolveRequest.getPods() == null) { - responder.sendStatus(HttpResponseStatus.BAD_REQUEST); - return; - } - - // Convert serialized pods back to Discoverable objects - List discoverables = new ArrayList<>(); - for (PodInfo podInfo : resolveRequest.getPods()) { - discoverables.add(new Discoverable("task.worker", - new InetSocketAddress(podInfo.getHost(), podInfo.getPort()))); - } - - Discoverable selectedPod = taskManager.resolvePod(resolveRequest.getNamespace(), discoverables); - - if (selectedPod == null) { - responder.sendStatus(HttpResponseStatus.SERVICE_UNAVAILABLE); - return; - } - - PodInfo responsePod = new PodInfo( - selectedPod.getSocketAddress().getHostString(), - selectedPod.getSocketAddress().getPort() - ); - - responder.sendJson(HttpResponseStatus.OK, GSON.toJson(responsePod)); - } catch (Exception e) { - LOG.error("Failed to resolve pod in Task Manager Service", e); - responder.sendString(HttpResponseStatus.INTERNAL_SERVER_ERROR, e.getMessage()); - } - } - - @POST - @Path("/finish") - public void finish(FullHttpRequest request, HttpResponder responder) { - try { - String jsonBody = request.content().toString(StandardCharsets.UTF_8); - FinishRequest finishRequest = GSON.fromJson(jsonBody, FinishRequest.class); - - if (finishRequest == null || finishRequest.getNamespace() == null || finishRequest.getPod() == null) { - responder.sendStatus(HttpResponseStatus.BAD_REQUEST); - return; - } - - Discoverable pod = new Discoverable("task.worker", - new InetSocketAddress(finishRequest.getPod().getHost(), finishRequest.getPod().getPort())); - - taskManager.finishTask(finishRequest.getNamespace(), pod, finishRequest.isRejected()); - responder.sendStatus(HttpResponseStatus.OK); - } catch (Exception e) { - LOG.error("Failed to finish task in Task Manager Service", e); - responder.sendString(HttpResponseStatus.INTERNAL_SERVER_ERROR, e.getMessage()); - } - } - - // DTO Classes for Serialization - public static class ResolveRequest { - private String namespace; - private List pods; - - public String getNamespace() { - return namespace; - } - - public List getPods() { - return pods; - } - } - - public static class FinishRequest { - private String namespace; - private PodInfo pod; - private boolean rejected; - - public String getNamespace() { - return namespace; - } - - public PodInfo getPod() { - return pod; - } - - public boolean isRejected() { - return rejected; - } - } - - public static class PodInfo { - private String host; - private int port; - - public PodInfo(String host, int port) { - this.host = host; - this.port = port; - } - - public String getHost() { - return host; - } - - public int getPort() { - return port; - } - } -} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java index c7cb33271a6e..99a267fcf292 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java @@ -26,8 +26,7 @@ public class TaskManagerServiceModule extends AbstractModule { @Override protected void configure() { - // Bind the core TaskManager as a singleton - bind(TaskManager.class).toProvider(TaskManager::getInstance).in(Scopes.SINGLETON); + // Bind the Netty Proxy service itself bind(TaskManagerService.class).in(Scopes.SINGLETON); diff --git a/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java b/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java deleted file mode 100644 index 6ca5893c417c..000000000000 --- a/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java +++ /dev/null @@ -1,115 +0,0 @@ -/* - * 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 java.net.InetSocketAddress; -import java.util.ArrayList; -import java.util.List; -import org.apache.twill.discovery.Discoverable; -import org.junit.Assert; -import org.junit.Before; -import org.junit.Test; - -/** - * Unit tests for {@link TaskManager} warm sticky lease orchestration. - */ -public class TaskManagerTest { - - private List pods; - - @Before - public void setUp() { - pods = new ArrayList<>(); - // Define 3 mock task worker pods - pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.1", 11015))); - pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.2", 11015))); - pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.3", 11015))); - } - - @Test - public void testStickyRoutingAndLeasing() { - TaskManager taskManager = TaskManager.getInstance(); - - // 1. First request for ns1 should claim a pod - Discoverable pod1 = taskManager.resolvePod("ns1", pods); - Assert.assertNotNull(pod1); - - // 2. Second request for ns1 should land on the same pod (Stickiness) - Discoverable pod1Repeat = taskManager.resolvePod("ns1", pods); - Assert.assertEquals(pod1.getSocketAddress(), pod1Repeat.getSocketAddress()); - - // 3. First request for ns2 should claim a different, idle pod (Isolation) - Discoverable pod2 = taskManager.resolvePod("ns2", pods); - Assert.assertNotNull(pod2); - Assert.assertNotEquals(pod1.getSocketAddress(), pod2.getSocketAddress()); - - // Clean up active tasks - taskManager.finishTask("ns1", pod1, false); - taskManager.finishTask("ns1", pod1Repeat, false); - taskManager.finishTask("ns2", pod2, false); - } - - @Test - public void testLogicalResetAfterMaxTasks() { - TaskManager taskManager = TaskManager.getInstance(); - - // 1. Claim a pod for ns3 - Discoverable initialPod = taskManager.resolvePod("ns3", pods); - Assert.assertNotNull(initialPod); - taskManager.finishTask("ns3", initialPod, false); - - // 2. Send tasks to reach the logical reset limit - for (int i = 0; i < TaskManager.MAX_TOTAL_TASKS_BEFORE_RESET - 1; i++) { - Discoverable p = taskManager.resolvePod("ns3", pods); - Assert.assertEquals(initialPod.getSocketAddress(), p.getSocketAddress()); - taskManager.finishTask("ns3", p, false); - } - - // 3. The next request for a DIFFERENT namespace (ns4) should now be able to claim this pod - // because it was logically reset (released) on the limit threshold! - Discoverable resetPod = taskManager.resolvePod("ns4", pods); - Assert.assertNotNull(resetPod); - Assert.assertEquals(initialPod.getSocketAddress(), resetPod.getSocketAddress()); - taskManager.finishTask("ns4", resetPod, false); - } - - @Test - public void testRejectionRevertsTotalProcessedCount() { - TaskManager taskManager = TaskManager.getInstance(); - - // 1. Claim a pod for ns5 - Discoverable initialPod = taskManager.resolvePod("ns5", pods); - Assert.assertNotNull(initialPod); - - // 2. Reject it -> totalProcessed should revert back to 0 - taskManager.finishTask("ns5", initialPod, true); - - // 3. Send 10 tasks to reach the logical reset limit - for (int i = 0; i < TaskManager.MAX_TOTAL_TASKS_BEFORE_RESET; i++) { - Discoverable p = taskManager.resolvePod("ns5", pods); - Assert.assertEquals(initialPod.getSocketAddress(), p.getSocketAddress()); - taskManager.finishTask("ns5", p, false); - } - - // 4. The next request for ns6 should claim this pod because it reset successfully after exactly 10 tasks, - // meaning the rejected task did not count towards the 10 task limit. - Discoverable resetPod = taskManager.resolvePod("ns6", pods); - Assert.assertNotNull(resetPod); - Assert.assertEquals(initialPod.getSocketAddress(), resetPod.getSocketAddress()); - taskManager.finishTask("ns6", resetPod, false); - } -} From f9045abfbe32ba2d1b2a1c5c67645e83d7562462 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 13:12:18 +0000 Subject: [PATCH 15/54] refactor: Add Predictive Proxy Expiration and Phase 0/1 Code Cleanup - Added timestamp tracking to PodState and ProxyBackendHandler to predict tenant lease expiration - Upgraded ProxyFrontendHandler 3-step hierarchy to safely route 35s expired Developer loads to enterprise idle pods - Deleted legacy TaskManager.java dead code - Replaced sidhirange with shruzard prefixes --- Dockerfile | 14 ++ _agents/rules/rbac_taskmanager.md | 18 ++ _agents/rules/task_manager_context.md | 19 ++ .../cdap/common/internal/remote/PodState.java | 56 +++++ .../internal/remote/ProxyBackendHandler.java | 109 +++++++++ .../internal/remote/ProxyFrontendHandler.java | 230 ++++++++++++++++++ .../common/internal/remote/RemoteClient.java | 163 +------------ .../internal/remote/RemoteTaskExecutor.java | 11 +- .../internal/remote/StickyLeaseManager.java | 25 +- .../common/internal/remote/TaskManager.java | 199 --------------- .../remote/TaskManagerHttpHandler.java | 157 ------------ .../internal/remote/TaskManagerService.java | 89 +++++-- .../remote/TaskManagerServiceModule.java | 8 +- .../remote/TaskWorkerHttpHandlerInternal.java | 49 ++-- cdap-ui | 2 +- task-manager-service.yaml | 16 +- 16 files changed, 590 insertions(+), 575 deletions(-) create mode 100644 Dockerfile create mode 100644 _agents/rules/rbac_taskmanager.md create mode 100644 _agents/rules/task_manager_context.md create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java create mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java delete mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java delete mode 100644 cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 000000000000..66e1fbb90ea8 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,14 @@ +FROM us-east1-docker.pkg.dev/cloud-data-fusion-images/cdf/cloud-data-fusion:latest +# For OSS CDAP, use "FROM gcr.io/cdapio/cdap:latest" + +RUN rm -rf /opt/cdap/master/ext/runtimeproviders \ + && rm -rf /opt/cdap/master/ext/runtimes \ + && rm -rf /opt/cdap/master/ext/environments \ + && rm -rf /opt/cdap/master/lib/io.cdap.cdap.cdap* \ + && rm -rf /opt/cdap/master/artifacts/spark3_2.12 + +COPY opt/cdap/master/lib /opt/cdap/master/lib +COPY opt/cdap/master/ext /opt/cdap/master/ext +COPY opt/cdap/master/artifacts/spark3_2.12/* /opt/cdap/master/artifacts/spark3_2.12/ + +RUN chmod -R 755 /opt/cdap diff --git a/_agents/rules/rbac_taskmanager.md b/_agents/rules/rbac_taskmanager.md new file mode 100644 index 000000000000..ea50b48a302a --- /dev/null +++ b/_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` diff --git a/_agents/rules/task_manager_context.md b/_agents/rules/task_manager_context.md new file mode 100644 index 000000000000..3055a1ae4a36 --- /dev/null +++ b/_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. diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java new file mode 100644 index 000000000000..20b3d1d2858b --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -0,0 +1,56 @@ +/* + * 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; + this.lastActivityTime = 0; // Instantly trigger predictions on boot + } + + 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; + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java new file mode 100644 index 000000000000..a9d6b77a7306 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -0,0 +1,109 @@ +/* + * 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 podRegistry; + private final String targetWorkerAddress; + + public ProxyBackendHandler(Channel inboundChannel, Map podRegistry, String targetWorkerAddress) { + this.inboundChannel = inboundChannel; + this.podRegistry = podRegistry; + this.targetWorkerAddress = targetWorkerAddress; + } + + @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"); + + if (activeTasksStr != null) { + try { + state.setInflightRequests(Integer.parseInt(activeTasksStr)); + } catch (NumberFormatException e) { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); + } + } else { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); + } + + + if (leasedNamespace != null) { + state.setLeasedNamespace(leasedNamespace); + } + + state.setLastActivityTime(System.currentTimeMillis()); + + if (activeTasksStr != null || leasedNamespace != null) { + LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState for {}. Occupancy: {}, Namespace: {}", + targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); + } + } + } + } + + // 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) { + ProxyFrontendHandler.closeOnFlush(inboundChannel); + } + + @Override + public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { + cause.printStackTrace(); + ProxyFrontendHandler.closeOnFlush(ctx.channel()); + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java new file mode 100644 index 000000000000..c00d7ddfa8d8 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -0,0 +1,230 @@ +/* + * 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.bootstrap.Bootstrap; +import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelFutureListener; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelInboundHandlerAdapter; +import io.netty.channel.ChannelInitializer; +import io.netty.channel.ChannelOption; +import io.netty.channel.ChannelPipeline; +import io.netty.channel.socket.SocketChannel; +import io.netty.channel.socket.nio.NioSocketChannel; +import io.netty.handler.codec.http.DefaultFullHttpResponse; +import io.netty.handler.codec.http.FullHttpResponse; +import io.netty.handler.codec.http.HttpClientCodec; +import io.netty.handler.codec.http.HttpContent; +import io.netty.handler.codec.http.HttpRequest; +import io.netty.handler.codec.http.HttpResponseStatus; +import io.netty.handler.codec.http.HttpVersion; +import io.netty.util.ReferenceCountUtil; + +import java.util.LinkedList; +import java.util.Map; +import java.util.Queue; +import java.util.Set; +import java.util.HashSet; +import org.apache.twill.discovery.Discoverable; +import org.apache.twill.discovery.DiscoveryServiceClient; +import io.cdap.cdap.common.conf.Constants; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { + + private static final Logger LOG = LoggerFactory.getLogger(ProxyFrontendHandler.class); + + + private final Map podRegistry; + private final DiscoveryServiceClient discoveryServiceClient; + private Channel outboundChannel; + private boolean connecting = false; + private final Queue pendingMessages = new LinkedList<>(); + + public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { + this.podRegistry = podRegistry; + this.discoveryServiceClient = discoveryServiceClient; + } + + @Override + public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { + if (msg instanceof HttpRequest) { + HttpRequest req = (HttpRequest) msg; + + // 0. Synchronous K8s Discovery (Zero-Stale State) + // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient evaluates a local memory cache backed by a push-based ZooKeeper watch. + Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); + Set activePods = new HashSet<>(); + for (Discoverable d : discoverables) { + activePods.add(d.getSocketAddress().getHostString() + ":" + d.getSocketAddress().getPort()); + } + + for (String podIp : activePods) { + podRegistry.putIfAbsent(podIp, new PodState(null, 0)); + } + podRegistry.keySet().removeIf(existingPod -> !activePods.contains(existingPod)); + + String targetNamespace = req.headers().get("X-CDF-Namespace"); + if (targetNamespace == null) targetNamespace = "default"; + + String targetWorkerAddress = null; + + // 1. Warm Match: Thread-safe scan specifically locking evaluation + for (Map.Entry entry : podRegistry.entrySet()) { + PodState state = entry.getValue(); + synchronized (state) { + if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { + targetWorkerAddress = entry.getKey(); + state.setInflightRequests(state.getInflightRequests() + 1); + LOG.info("shruzard - ProxyFrontendHandler: Found warm match for '{}' at {}. Occupancy: {}", targetNamespace, targetWorkerAddress, state.getInflightRequests()); + break; + } + } + } + + // 2. Idle Choice: Thread-safe claim of an unleased pod, OR an expired pod (35s predicted timeout avoiding clock drift) + if (targetWorkerAddress == null) { + for (Map.Entry entry : podRegistry.entrySet()) { + PodState state = entry.getValue(); + synchronized (state) { + boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); + boolean isExpiredIdle = (state.getInflightRequests() == 0 && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); + + if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { + targetWorkerAddress = entry.getKey(); + state.setLeasedNamespace(targetNamespace); + state.setInflightRequests(1); + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod (Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", + isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); + break; + } + } + } + } + + // 3. Busy Rejection: All pods saturated + if (targetWorkerAddress == null) { + LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased incorrectly. Rejecting request for namespace '{}'", targetNamespace); + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + ReferenceCountUtil.release(msg); + return; + } + + final String chosenWorker = targetWorkerAddress; + String[] hostPort = targetWorkerAddress.split(":"); + + // Apply backpressure on client until connection established + ctx.channel().config().setAutoRead(false); + connecting = true; + + Bootstrap b = new Bootstrap(); + b.group(ctx.channel().eventLoop()) + .channel(NioSocketChannel.class) + .option(ChannelOption.SO_KEEPALIVE, true) + .handler(new ChannelInitializer() { + @Override + protected void initChannel(SocketChannel ch) { + ChannelPipeline p = ch.pipeline(); + p.addLast(new HttpClientCodec()); + p.addLast(new ProxyBackendHandler(ctx.channel(), podRegistry, chosenWorker)); + } + }); + + ChannelFuture f = b.connect(hostPort[0], Integer.parseInt(hostPort[1])); + outboundChannel = f.channel(); + + f.addListener((ChannelFutureListener) future -> { + connecting = false; + if (future.isSuccess()) { + Object pendingMsg = pendingMessages.poll(); + while (pendingMsg != null) { + outboundChannel.write(pendingMsg); + pendingMsg = pendingMessages.poll(); + } + outboundChannel.flush(); + ctx.channel().config().setAutoRead(true); + } else { + Object pendingMsg = pendingMessages.poll(); + while (pendingMsg != null) { + ReferenceCountUtil.release(pendingMsg); + pendingMsg = pendingMessages.poll(); + } + // Thread-safe decrement on fallback + PodState fallbackState = podRegistry.get(chosenWorker); + if (fallbackState != null) { + synchronized (fallbackState) { + fallbackState.setInflightRequests(Math.max(0, fallbackState.getInflightRequests() - 1)); + } + } + ctx.channel().close(); + } + }); + + pendingMessages.add(ReferenceCountUtil.retain(msg)); + + } else if (msg instanceof HttpContent) { + if (connecting) { + pendingMessages.add(ReferenceCountUtil.retain(msg)); + } else if (outboundChannel != null && outboundChannel.isActive()) { + outboundChannel.writeAndFlush(ReferenceCountUtil.retain(msg)); + } else { + ReferenceCountUtil.release(msg); + } + } + } + + @Override + public void channelReadComplete(ChannelHandlerContext ctx) { + if (outboundChannel != null && outboundChannel.isActive() && !connecting) { + outboundChannel.flush(); + } + ctx.fireChannelReadComplete(); + } + + @Override + public void channelWritabilityChanged(ChannelHandlerContext ctx) { + if (outboundChannel != null && outboundChannel.isActive()) { + outboundChannel.config().setAutoRead(ctx.channel().isWritable()); + } + ctx.fireChannelWritabilityChanged(); + } + + @Override + public void channelInactive(ChannelHandlerContext ctx) { + if (outboundChannel != null) { + closeOnFlush(outboundChannel); + } + } + + @Override + public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { + cause.printStackTrace(); + closeOnFlush(ctx.channel()); + } + + static void closeOnFlush(Channel ch) { + if (ch.isActive()) { + ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE); + } + } +} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index 9bea26792ea0..3f93f327f5a0 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -67,8 +67,7 @@ public class RemoteClient { public static final String RUNTIME_SERVICE_ROUTING_BASE_URI = "cdap.runtime.service.routing.base.uri"; private static final Logger LOG = LoggerFactory.getLogger(RemoteClient.class); - private static final ThreadLocal CURRENT_RESOLVED_POD = new ThreadLocal<>(); - private static final ThreadLocal CURRENT_ROUTING_KEY = new ThreadLocal<>(); + private static final String TASK_MANAGER_URL = "http://cdap-task-manager.default.svc.cluster.local:11025"; private static final Gson GSON = new Gson(); @@ -181,13 +180,9 @@ private HttpResponse executeNonIdempotent(HttpRequest request) throws IOExceptio HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength()); - boolean rejected = false; try { HttpResponse response = HttpRequests.execute(httpRequest, httpRequestConfig); int responseCode = response.getResponseCode(); - if (responseCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - rejected = true; - } // 503 is always retryable. Other 5xx errors are retryable if the request is idempotent (handled in // RemoteClient#executeIdempotent(HttpRequest) if (responseCode == HttpURLConnection.HTTP_UNAVAILABLE) { @@ -215,19 +210,7 @@ private HttpResponse executeNonIdempotent(HttpRequest request) throws IOExceptio } return response; } catch (ConnectException e) { - rejected = true; throw new ServiceUnavailableException(discoverableServiceName, e); - } catch (IOException | RuntimeException e) { - rejected = true; - throw e; - } finally { - Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); - String routingKey = CURRENT_ROUTING_KEY.get(); - if (resolvedPod != null && routingKey != null) { - notifyTaskManagerFinished(routingKey, resolvedPod, rejected); - } - CURRENT_RESOLVED_POD.remove(); - CURRENT_ROUTING_KEY.remove(); } } @@ -242,59 +225,15 @@ public void executeStreamingRequest(HttpRequest request) HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength(), request.getConsumer()); - boolean rejected = false; try { HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { - if (httpResponse.getResponseCode() == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - rejected = true; - } throw new IOException( String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), httpResponse.getResponseCode())); } httpResponse.consumeContent(); - } catch (IOException | RuntimeException e) { - rejected = true; - throw e; - } finally { - Discoverable resolvedPod = CURRENT_RESOLVED_POD.get(); - String routingKey = CURRENT_ROUTING_KEY.get(); - if (resolvedPod != null && routingKey != null) { - notifyTaskManagerFinished(routingKey, resolvedPod, rejected); - } - CURRENT_RESOLVED_POD.remove(); - CURRENT_ROUTING_KEY.remove(); - } - } - private void notifyTaskManagerFinished(String namespace, Discoverable pod, boolean rejected) { - try { - URL url = new URL(TASK_MANAGER_URL + "/v3/taskmanager/finish"); - TaskManagerHttpHandler.FinishRequest finishRequest = new TaskManagerHttpHandler.FinishRequest(); - - java.lang.reflect.Field nsField = finishRequest.getClass().getDeclaredField("namespace"); - nsField.setAccessible(true); - nsField.set(finishRequest, namespace); - - TaskManagerHttpHandler.PodInfo podInfo = new TaskManagerHttpHandler.PodInfo( - pod.getSocketAddress().getHostString(), pod.getSocketAddress().getPort()); - java.lang.reflect.Field podField = finishRequest.getClass().getDeclaredField("pod"); - podField.setAccessible(true); - podField.set(finishRequest, podInfo); - - java.lang.reflect.Field rejectedField = finishRequest.getClass().getDeclaredField("rejected"); - rejectedField.setAccessible(true); - rejectedField.set(finishRequest, rejected); - - HttpRequest req = HttpRequest.post(url) - .addHeader(HttpHeaders.CONTENT_TYPE, "application/json") - .withBody(GSON.toJson(finishRequest)) - .build(); - - HttpRequests.execute(req, httpRequestConfig); - } catch (Exception e) { - LOG.warn("sidhdirenge - Failed to notify Task Manager of task completion", e); } } @@ -368,104 +307,12 @@ public URL resolve(String resource, @Nullable String routingKey) { } } - LOG.info("sidhdirenge - RemoteClient resolving stickily via TaskManager for service {} with routingKey: {}", - discoverableServiceName, routingKey); - - // 1. Fetch all currently discovered endpoints - Iterable discoverables = () -> discoveryClient.discover(discoverableServiceName) - .iterator(); - List list = new ArrayList<>(); - for (Discoverable d : discoverables) { - // Perform DNS lookup to resolve the service hostname into individual pod IPs (for headless services) - try { - java.net.InetAddress[] addresses = java.net.InetAddress.getAllByName( - d.getSocketAddress().getHostName()); - for (java.net.InetAddress addr : addresses) { - list.add(new Discoverable(d.getName(), - new java.net.InetSocketAddress(addr.getHostAddress(), d.getSocketAddress().getPort()), - d.getPayload())); - } - } catch (java.net.UnknownHostException e) { - // Fallback to original discoverable if DNS lookup fails - list.add(d); - } - } - - if (list.isEmpty()) { - throw new ServiceUnavailableException(discoverableServiceName); - } - - // 2. Sort endpoints by IP address and port to ensure consistent ordering across all client instances - list.sort(Comparator.comparing((Discoverable d) -> d.getSocketAddress().getHostName()) - .thenComparingInt(d -> d.getSocketAddress().getPort())); - - // 3. Delegate to the standalone TaskManager Service over HTTP - Discoverable discoverable = null; - try { - URL url = new URL(TASK_MANAGER_URL + "/v3/taskmanager/resolve"); - TaskManagerHttpHandler.ResolveRequest resolveRequest = new TaskManagerHttpHandler.ResolveRequest(); - - java.lang.reflect.Field nsField = resolveRequest.getClass().getDeclaredField("namespace"); - nsField.setAccessible(true); - nsField.set(resolveRequest, routingKey); - - List podInfos = new ArrayList<>(); - for (Discoverable pod : list) { - podInfos.add(new TaskManagerHttpHandler.PodInfo( - pod.getSocketAddress().getHostString(), pod.getSocketAddress().getPort())); - } - java.lang.reflect.Field podsField = resolveRequest.getClass().getDeclaredField("pods"); - podsField.setAccessible(true); - podsField.set(resolveRequest, podInfos); - - HttpRequest req = HttpRequest.post(url) - .addHeader(HttpHeaders.CONTENT_TYPE, "application/json") - .withBody(GSON.toJson(resolveRequest)) - .build(); - - HttpResponse resp = HttpRequests.execute(req, httpRequestConfig); - if (resp.getResponseCode() == HttpURLConnection.HTTP_OK) { - TaskManagerHttpHandler.PodInfo selectedPodInfo = GSON.fromJson( - resp.getResponseBodyAsString(), TaskManagerHttpHandler.PodInfo.class); - - byte[] payload = list.isEmpty() ? new byte[0] : list.get(0).getPayload(); - for (Discoverable d : list) { - if (d.getSocketAddress().getPort() == selectedPodInfo.getPort() - && (d.getSocketAddress().getHostName().equals(selectedPodInfo.getHost()) - || (d.getSocketAddress().getAddress() != null - && d.getSocketAddress().getAddress().getHostAddress().equals(selectedPodInfo.getHost())))) { - payload = d.getPayload(); - break; - } - } - discoverable = new Discoverable("task.worker", - new java.net.InetSocketAddress(selectedPodInfo.getHost(), selectedPodInfo.getPort()), payload); - } - } catch (Exception e) { - LOG.warn("sidhdirenge - Failed to resolve pod via Task Manager HTTP Service. Falling back to local hashing.", e); - } - - // Fallback: If Task Manager is down or returns error, use standard consistent hashing - if (discoverable == null) { - int baseIndex = (routingKey.hashCode() & Integer.MAX_VALUE) % list.size(); - discoverable = list.get(baseIndex); - LOG.warn("sidhdirenge - TaskManager resolution failed. Falling back to default index {}", baseIndex); - } - - // Store resolved pod context in ThreadLocal for task execution callbacks - CURRENT_RESOLVED_POD.set(discoverable); - CURRENT_ROUTING_KEY.set(routingKey); - - LOG.info("sidhdirenge - Centralized TaskManager selected warm pod IP {} for routingKey: {}", - discoverable.getSocketAddress(), routingKey); - - URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); + LOG.info("shruzard - RemoteClient routing directly to Netty TaskManager L7 proxy for routingKey: {}", routingKey); try { - return rewriteUrl(uri.toURL()); + String cleanPath = (basePath + resource).replaceAll("//+", "/"); + return new URL(TASK_MANAGER_URL + "/" + cleanPath); } catch (MalformedURLException e) { - throw new IllegalStateException( - String.format("Discovered service %s, but it announced malformed URL %s", - discoverableServiceName, uri), e); + throw new ServiceUnavailableException(discoverableServiceName, e); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index aa4682c8fec2..9847d8ffde64 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -139,12 +139,19 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception String embeddedNamespace = runnableTaskRequest.getParam().getEmbeddedTaskRequest().getNamespace(); if (embeddedNamespace != null && !embeddedNamespace.isEmpty()) { namespace = embeddedNamespace; - LOG.info("sidhdirenge - RemoteTaskExecutor: Mapped SystemAppTask namespace to embedded: {}", + LOG.info("shruzard - RemoteTaskExecutor: Mapped SystemAppTask namespace to embedded: {}", namespace); } } + String routingKey = namespace; + if (System.currentTimeMillis() - startTime > 60000) { + LOG.warn("shruzard - TaskManager Proxy unreachable for 60s! Bypassing proxy and falling back to direct Worker routing!"); + routingKey = null; // Setting to null triggers CDAP's native RandomEndpoint discovery in RemoteClient + } + HttpRequest.Builder requestBuilder = remoteClient - .requestBuilder(HttpMethod.POST, workerUrl, namespace) + .requestBuilder(HttpMethod.POST, workerUrl, routingKey) + .addHeader("X-CDF-Namespace", namespace) .withBody(requestBody.duplicate()); if (compression) { requestBuilder.addHeader(HttpHeaders.CONTENT_ENCODING, "gzip"); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index 12754c57c642..6002725b8ad3 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -20,6 +20,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import javax.annotation.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,15 +45,21 @@ public class StickyLeaseManager { private final AtomicInteger activeTaskCount = new AtomicInteger(0); private final AtomicInteger totalTasksProcessedInLease = new AtomicInteger(0); private volatile long lastActivityTimeMillis; + private final Consumer onLeaseAcquired; + private final Runnable onLeaseReleased; public StickyLeaseManager() { - this(10, 10); + this(10, 10, null, null); } - public StickyLeaseManager(int maxConcurrentTasks, int maxTasksPerLease) { + public StickyLeaseManager(int maxConcurrentTasks, int maxTasksPerLease, + Consumer onLeaseAcquired, + Runnable onLeaseReleased) { this.maxConcurrentTasks = maxConcurrentTasks; this.maxTasksPerLease = maxTasksPerLease; this.lastActivityTimeMillis = System.currentTimeMillis(); + this.onLeaseAcquired = onLeaseAcquired; + this.onLeaseReleased = onLeaseReleased; } /** @@ -71,8 +78,11 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant lastActivityTimeMillis = System.currentTimeMillis(); long elapsed = System.currentTimeMillis() - claimStartTime; LOG.info( - "Lease claimed by namespace '{}' (Tier: {}) in {}ms (Boot penalty entirely avoided)", + "shruzard Lease claimed by namespace '{}' (Tier: {}) in {}ms (Boot penalty entirely avoided)", namespace.getNamespace(), tier, elapsed); + if (onLeaseAcquired != null) { + onLeaseAcquired.accept(namespace); + } return AcquisitionStatus.SUCCESS; } } @@ -84,7 +94,7 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant } // Mismatching namespace -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) - LOG.info("Enforcement: Rejecting request for namespace '{}', current lease is held by '{}'", + LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}'", namespace.getNamespace(), currentLease.get()); return AcquisitionStatus.REJECTED_MISMATCH; } @@ -99,7 +109,7 @@ public synchronized AcquisitionStatus startTask(NamespaceId namespace, TenantTie } if (activeTaskCount.get() >= maxConcurrentTasks) { - LOG.info("Concurrency limit reached ({} tasks active) for namespace '{}'", + LOG.info("shruzard - StickyLeaseManager: Concurrency limit reached ({} tasks active) for namespace '{}'", activeTaskCount.get(), namespace.getNamespace()); return AcquisitionStatus.REJECTED_MAX_CONCURRENCY; } @@ -153,8 +163,11 @@ public synchronized void releaseLease(String reason) { if (oldNamespace != null) { activeTaskCount.set(0); totalTasksProcessedInLease.set(0); - LOG.info("Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", + LOG.info("shruzard - StickyLeaseManager: Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", oldNamespace.getNamespace(), reason); + if (onLeaseReleased != null) { + onLeaseReleased.run(); + } } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java deleted file mode 100644 index 0ab82fe9967f..000000000000 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManager.java +++ /dev/null @@ -1,199 +0,0 @@ -/* - * 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 java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.locks.ReentrantLock; -import javax.annotation.Nullable; -import org.apache.twill.discovery.Discoverable; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Centralized Task Manager for orchestrating Warm Sticky Leases on Task Worker pods. - * This class coordinates leases, concurrency, and logical resets. - */ -public class TaskManager { - - private static final Logger LOG = LoggerFactory.getLogger(TaskManager.class); - private static final TaskManager INSTANCE = new TaskManager(); - - // Concurrency and task limits based on the design doc - private static final int MAX_CONCURRENT_TASKS_PER_POD = 10; - static final int MAX_TOTAL_TASKS_BEFORE_RESET = 10; - - private final ReentrantLock lock = new ReentrantLock(); - - // Lease state maps - // Pod IP/Key -> Namespace currently leased - private final Map podLeases = new HashMap<>(); - // Pod IP/Key -> Active concurrent task count - private final Map podActiveTaskCounts = new HashMap<>(); - // Pod IP/Key -> Total tasks processed on the current lease - private final Map podTotalTaskProcessedCounts = new HashMap<>(); - - public static TaskManager getInstance() { - return INSTANCE; - } - - private TaskManager() { - // Singleton - } - - /** - * Resolves the target warm pod for a given namespace based on the sticky lease model. - * - * @param namespace the namespace requesting execution - * @param availablePods the list of currently discovered pods - * @return the selected pod, or null if no pod is available - */ - @Nullable - public Discoverable resolvePod(String namespace, List availablePods) { - lock.lock(); - try { - String leasedPodIp = null; - Discoverable selectedPod = null; - - // 1. Find if a pod is already leased to this namespace and has capacity - for (Discoverable pod : availablePods) { - String podIp = getPodKey(pod); - String currentLease = podLeases.get(podIp); - - if (namespace.equals(currentLease)) { - int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); - if (activeTasks < MAX_CONCURRENT_TASKS_PER_POD && totalProcessed < MAX_TOTAL_TASKS_BEFORE_RESET) { - leasedPodIp = podIp; - selectedPod = pod; - break; - } - } - } - - // 2. If no active lease exists (or it is at capacity), find an idle/unleased pod - if (selectedPod == null) { - for (Discoverable pod : availablePods) { - String podIp = getPodKey(pod); - String currentLease = podLeases.get(podIp); - - if (currentLease == null) { - // Establish a new lease on this idle pod - podLeases.put(podIp, namespace); - podActiveTaskCounts.put(podIp, 0); - podTotalTaskProcessedCounts.put(podIp, 0); - - LOG.info("sidhdirenge - TaskManager: Established new lease for namespace '{}' on pod '{}'", - namespace, podIp); - - leasedPodIp = podIp; - selectedPod = pod; - break; - } - } - } - - // 3. Fallback: If all pods are leased to other namespaces, find the pod with the least load - if (selectedPod == null) { - LOG.warn("sidhdirenge - TaskManager: All pods are leased. Falling back to least-loaded pod."); - int minLoad = Integer.MAX_VALUE; - List bestPods = new ArrayList<>(); - for (Discoverable pod : availablePods) { - String podIp = getPodKey(pod); - int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); - if (activeTasks < minLoad) { - minLoad = activeTasks; - bestPods.clear(); - bestPods.add(pod); - } else if (activeTasks == minLoad) { - bestPods.add(pod); - } - } - - if (!bestPods.isEmpty()) { - int randomIndex = java.util.concurrent.ThreadLocalRandom.current().nextInt(bestPods.size()); - selectedPod = bestPods.get(randomIndex); - leasedPodIp = getPodKey(selectedPod); - - // Force-assign lease to the new namespace - podLeases.put(leasedPodIp, namespace); - podActiveTaskCounts.put(leasedPodIp, 0); - podTotalTaskProcessedCounts.put(leasedPodIp, 0); - } - } - - // 4. Increment task counts for the selected pod - if (selectedPod != null) { - int activeTasks = podActiveTaskCounts.getOrDefault(leasedPodIp, 0) + 1; - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(leasedPodIp, 0) + 1; - - podActiveTaskCounts.put(leasedPodIp, activeTasks); - podTotalTaskProcessedCounts.put(leasedPodIp, totalProcessed); - - LOG.info("sidhdirenge - TaskManager: Routing task for '{}' to pod '{}' (Active: {}, Total: {})", - namespace, leasedPodIp, activeTasks, totalProcessed); - } - - return selectedPod; - } finally { - lock.unlock(); - } - } - - public void finishTask(String namespace, Discoverable pod, boolean rejected) { - lock.lock(); - try { - String podIp = getPodKey(pod); - int activeTasks = podActiveTaskCounts.getOrDefault(podIp, 0); - if (activeTasks > 0) { - activeTasks = activeTasks - 1; - podActiveTaskCounts.put(podIp, activeTasks); - } - - if (rejected) { - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); - if (totalProcessed > 0) { - podTotalTaskProcessedCounts.put(podIp, totalProcessed - 1); - } - } else { - int totalProcessed = podTotalTaskProcessedCounts.getOrDefault(podIp, 0); - if (totalProcessed >= MAX_TOTAL_TASKS_BEFORE_RESET && activeTasks == 0) { - LOG.info("sidhdirenge - TaskManager: Pod '{}' finished all active tasks " - + "after reaching reset threshold. Reclaiming lease.", podIp); - releaseLease(podIp); - } - } - - LOG.info("sidhdirenge - TaskManager: Task finished for '{}' on pod '{}' (Remaining active: {}, Rejected: {})", - namespace, podIp, activeTasks, rejected); - } finally { - lock.unlock(); - } - } - - private void releaseLease(String podIp) { - podLeases.remove(podIp); - podActiveTaskCounts.remove(podIp); - podTotalTaskProcessedCounts.remove(podIp); - } - - private String getPodKey(Discoverable pod) { - return pod.getSocketAddress().getHostString() + ":" + pod.getSocketAddress().getPort(); - } -} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java deleted file mode 100644 index d53ea147929b..000000000000 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerHttpHandler.java +++ /dev/null @@ -1,157 +0,0 @@ -/* - * 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 com.google.gson.Gson; -import com.google.gson.reflect.TypeToken; -import io.cdap.http.AbstractHttpHandler; -import io.cdap.http.HttpResponder; -import io.netty.handler.codec.http.FullHttpRequest; -import io.netty.handler.codec.http.HttpResponseStatus; -import org.apache.twill.discovery.Discoverable; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.lang.reflect.Type; -import java.net.InetSocketAddress; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.List; -import javax.ws.rs.POST; -import javax.ws.rs.Path; - -/** - * Netty HTTP Handler for the standalone Task Manager Service. - */ -@Path("/v3/taskmanager") -public class TaskManagerHttpHandler extends AbstractHttpHandler { - - private static final Logger LOG = LoggerFactory.getLogger(TaskManagerHttpHandler.class); - private static final Gson GSON = new Gson(); - private final TaskManager taskManager = TaskManager.getInstance(); - - @POST - @Path("/resolve") - public void resolve(FullHttpRequest request, HttpResponder responder) { - try { - String jsonBody = request.content().toString(StandardCharsets.UTF_8); - ResolveRequest resolveRequest = GSON.fromJson(jsonBody, ResolveRequest.class); - - if (resolveRequest == null || resolveRequest.getNamespace() == null || resolveRequest.getPods() == null) { - responder.sendStatus(HttpResponseStatus.BAD_REQUEST); - return; - } - - // Convert serialized pods back to Discoverable objects - List discoverables = new ArrayList<>(); - for (PodInfo podInfo : resolveRequest.getPods()) { - discoverables.add(new Discoverable("task.worker", - new InetSocketAddress(podInfo.getHost(), podInfo.getPort()))); - } - - Discoverable selectedPod = taskManager.resolvePod(resolveRequest.getNamespace(), discoverables); - - if (selectedPod == null) { - responder.sendStatus(HttpResponseStatus.SERVICE_UNAVAILABLE); - return; - } - - PodInfo responsePod = new PodInfo( - selectedPod.getSocketAddress().getHostString(), - selectedPod.getSocketAddress().getPort() - ); - - responder.sendJson(HttpResponseStatus.OK, GSON.toJson(responsePod)); - } catch (Exception e) { - LOG.error("Failed to resolve pod in Task Manager Service", e); - responder.sendString(HttpResponseStatus.INTERNAL_SERVER_ERROR, e.getMessage()); - } - } - - @POST - @Path("/finish") - public void finish(FullHttpRequest request, HttpResponder responder) { - try { - String jsonBody = request.content().toString(StandardCharsets.UTF_8); - FinishRequest finishRequest = GSON.fromJson(jsonBody, FinishRequest.class); - - if (finishRequest == null || finishRequest.getNamespace() == null || finishRequest.getPod() == null) { - responder.sendStatus(HttpResponseStatus.BAD_REQUEST); - return; - } - - Discoverable pod = new Discoverable("task.worker", - new InetSocketAddress(finishRequest.getPod().getHost(), finishRequest.getPod().getPort())); - - taskManager.finishTask(finishRequest.getNamespace(), pod, finishRequest.isRejected()); - responder.sendStatus(HttpResponseStatus.OK); - } catch (Exception e) { - LOG.error("Failed to finish task in Task Manager Service", e); - responder.sendString(HttpResponseStatus.INTERNAL_SERVER_ERROR, e.getMessage()); - } - } - - // DTO Classes for Serialization - public static class ResolveRequest { - private String namespace; - private List pods; - - public String getNamespace() { - return namespace; - } - - public List getPods() { - return pods; - } - } - - public static class FinishRequest { - private String namespace; - private PodInfo pod; - private boolean rejected; - - public String getNamespace() { - return namespace; - } - - public PodInfo getPod() { - return pod; - } - - public boolean isRejected() { - return rejected; - } - } - - public static class PodInfo { - private String host; - private int port; - - public PodInfo(String host, int port) { - this.host = host; - this.port = port; - } - - public String getHost() { - return host; - } - - public int getPort() { - return port; - } - } -} diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 37510d9d7d3f..346605e8d9ae 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -19,49 +19,90 @@ import com.google.common.util.concurrent.AbstractIdleService; import com.google.inject.Inject; import io.cdap.cdap.common.conf.CConfiguration; -import io.cdap.cdap.common.http.CommonNettyHttpServiceFactory; -import io.cdap.http.NettyHttpService; +import io.netty.bootstrap.ServerBootstrap; +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelInitializer; +import io.netty.channel.ChannelPipeline; +import io.netty.channel.EventLoopGroup; +import io.netty.channel.nio.NioEventLoopGroup; +import io.netty.channel.socket.SocketChannel; +import io.netty.channel.socket.nio.NioServerSocketChannel; +import io.netty.handler.codec.http.HttpServerCodec; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Collections; +import org.apache.twill.discovery.Discoverable; +import org.apache.twill.discovery.DiscoveryServiceClient; +import io.cdap.cdap.common.conf.Constants; +import java.util.Set; +import java.util.HashSet; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; /** - * Guice-managed service that runs the Centralized Task Manager HTTP Server. + * Guice-managed service that runs the Centralized Task Manager HTTP Server (Netty Proxy POC). */ public class TaskManagerService extends AbstractIdleService { private static final Logger LOG = LoggerFactory.getLogger(TaskManagerService.class); - private final NettyHttpService httpService; - @Inject - TaskManagerService(CConfiguration cConf, - CommonNettyHttpServiceFactory commonNettyHttpServiceFactory, - TaskManagerHttpHandler taskManagerHttpHandler) { - - int port = cConf.getInt("task.manager.port", 11025); - String address = cConf.get("task.manager.address", "0.0.0.0"); + private final int port; + private final String address; + private EventLoopGroup bossGroup; + private EventLoopGroup workerGroup; + private ChannelFuture channelFuture; + + private final Map podRegistry = new ConcurrentHashMap<>(); - LOG.info("sidhdirenge - Initializing TaskManagerService on {}:{}", address, port); + private final DiscoveryServiceClient discoveryServiceClient; + + @Inject + TaskManagerService(CConfiguration cConf, DiscoveryServiceClient discoveryServiceClient) { + this.port = cConf.getInt("task.manager.port", 11025); + this.address = cConf.get("task.manager.address", "0.0.0.0"); + this.discoveryServiceClient = discoveryServiceClient; - this.httpService = commonNettyHttpServiceFactory.builder("task-manager", false) - .setHost(address) - .setPort(port) - .setHttpHandlers(Collections.singletonList(taskManagerHttpHandler)) - .build(); + LOG.info("shruzard - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); } @Override protected void startUp() throws Exception { - LOG.info("sidhdirenge - Starting TaskManagerService HTTP server..."); - httpService.start(); - LOG.info("sidhdirenge - TaskManagerService HTTP server started successfully at {}", httpService.getBindAddress()); + LOG.info("shruzard - Starting TaskManagerService Proxy HTTP server..."); + + bossGroup = new NioEventLoopGroup(1, + new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-boss-thread-%d").build()); + workerGroup = new NioEventLoopGroup(0, + new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-worker-thread-%d").build()); + + ServerBootstrap b = new ServerBootstrap(); + b.group(bossGroup, workerGroup) + .channel(NioServerSocketChannel.class) + .childHandler(new ChannelInitializer() { + @Override + protected void initChannel(SocketChannel ch) { + ChannelPipeline p = ch.pipeline(); + p.addLast(new HttpServerCodec()); + // NOTICE: NO HttpObjectAggregator here! + p.addLast(new ProxyFrontendHandler(podRegistry, discoveryServiceClient)); + } + }); + + channelFuture = b.bind(address, port).sync(); + LOG.info("shruzard - TaskManagerService Proxy HTTP server started successfully at {}:{}", address, port); } @Override protected void shutDown() throws Exception { - LOG.info("sidhdirenge - Stopping TaskManagerService HTTP server..."); - httpService.stop(); - LOG.info("sidhdirenge - TaskManagerService HTTP server stopped."); + LOG.info("shruzard - Stopping TaskManagerService Proxy HTTP server..."); + if (channelFuture != null) { + channelFuture.channel().close().sync(); + } + if (bossGroup != null) { + bossGroup.shutdownGracefully(); + } + if (workerGroup != null) { + workerGroup.shutdownGracefully(); + } + LOG.info("shruzard - TaskManagerService Proxy HTTP server stopped."); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java index fb470462ab7e..99a267fcf292 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java @@ -26,13 +26,9 @@ public class TaskManagerServiceModule extends AbstractModule { @Override protected void configure() { - // Bind the core TaskManager as a singleton - bind(TaskManager.class).toProvider(TaskManager::getInstance).in(Scopes.SINGLETON); - - // Bind the HTTP handler - bind(TaskManagerHttpHandler.class).in(Scopes.SINGLETON); + - // Bind the service itself + // Bind the Netty Proxy service itself bind(TaskManagerService.class).in(Scopes.SINGLETON); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java index b7726c46af23..24666e4d3044 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java @@ -121,7 +121,24 @@ public TaskWorkerHttpHandlerInternal(CConfiguration cConf, TaskWorker.USER_CODE_ISOLATION_ENABLED); this.concurrentRequestLimit = cConf.getInt(TaskWorker.REQUEST_LIMIT); int maxTasksPerLease = cConf.getInt("task.worker.lease.max.tasks", 10); - this.stickyLeaseManager = new StickyLeaseManager(concurrentRequestLimit, maxTasksPerLease); + this.stickyLeaseManager = new StickyLeaseManager(concurrentRequestLimit, maxTasksPerLease, + (namespaceId) -> { + try { + GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, cConf); + } catch (Exception e) { + LOG.warn("Failed to set GCP metadata task context for namespace {}", namespaceId, e); + } + }, + () -> { + try { + // Wipe Sidecar IAM Tokens (Identity / Security boundary) + GcpMetadataTaskContextUtil.clearGcpMetadataTaskContext(cConf); + // TODO: Wipe the JVM artifactCache to enforce disk reclaim boundary here + } catch (Exception e) { + LOG.warn("Failed to clear GCP metadata task context", e); + } + } + ); ScheduledExecutorService leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( Threads.createDaemonThreadFactory("lease-reclamation")); @@ -228,7 +245,7 @@ private void stopAndShutdown(ScheduledExecutorService executorService, Consumer< @POST @Path("/run") public void run(FullHttpRequest request, HttpResponder responder) { - LOG.info("sidhdirenge - Received task on worker {} for namespace :{}", + LOG.info("shruzard - Received task on worker {} for namespace :{}", System.getenv("HOSTNAME") != null ? System.getenv("HOSTNAME") : "unknown", request.headers()); if (mustRestart.get()) { @@ -267,15 +284,21 @@ public void run(FullHttpRequest request, HttpResponder responder) { if (leaseStatus != StickyLeaseManager.AcquisitionStatus.SUCCESS) { LOG.warn("Rejecting request for namespace {} due to lease status: {}", namespaceId, leaseStatus); - responder.sendStatus(HttpResponseStatus.TOO_MANY_REQUESTS); + + String currentLease = stickyLeaseManager.getCurrentLease() != null ? + stickyLeaseManager.getCurrentLease().getNamespace() : ""; + + responder.sendString(HttpResponseStatus.TOO_MANY_REQUESTS, + "Rejected due to lease status: " + leaseStatus.name(), + new DefaultHttpHeaders() + .add("X-Active-Tasks", String.valueOf(stickyLeaseManager.getActiveTaskCount())) + .add("X-Leased-Namespace", currentLease)); return; } runningRequestCount.incrementAndGet(); try { - // set the GcpMetadataTaskContext before running the task. - GcpMetadataTaskContextUtil.setGcpMetadataTaskContext(namespaceId, cConf); RunnableTaskContext runnableTaskContext = new RunnableTaskContext(runnableTaskRequest); runnableTaskLauncher.launchRunnableTask(runnableTaskContext); @@ -283,11 +306,16 @@ public void run(FullHttpRequest request, HttpResponder responder) { startTime, runnableTaskContext.isTerminateOnComplete(), runnableTaskRequest); + String currentLease = stickyLeaseManager.getCurrentLease() != null ? + stickyLeaseManager.getCurrentLease().getNamespace() : ""; + responder.sendContent(HttpResponseStatus.OK, new RunnableTaskBodyProducer(runnableTaskContext, taskCompletionConsumer, taskDetails), - new DefaultHttpHeaders().add(HttpHeaders.CONTENT_TYPE, - MediaType.APPLICATION_OCTET_STREAM)); + new DefaultHttpHeaders() + .add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_OCTET_STREAM) + .add("X-Active-Tasks", String.valueOf(stickyLeaseManager.getActiveTaskCount())) + .add("X-Leased-Namespace", currentLease)); } catch (ClassNotFoundException | ClassCastException ex) { responder.sendString(HttpResponseStatus.BAD_REQUEST, exceptionToJson(ex), @@ -307,13 +335,6 @@ public void run(FullHttpRequest request, HttpResponder responder) { // Potentially ran user code, hence terminate the runner. taskCompletionConsumer.accept(false, new TaskDetails(metricsCollectionService, startTime, true, runnableTaskRequest)); - } finally { - // clear the GcpMetadataTaskContext after the task is completed. - try { - GcpMetadataTaskContextUtil.clearGcpMetadataTaskContext(cConf); - } catch (Exception e) { - LOG.warn("Failed to clear GCP metadata task context", e); - } } } diff --git a/cdap-ui b/cdap-ui index 46e0c36dac1c..9d55597f75a9 160000 --- a/cdap-ui +++ b/cdap-ui @@ -1 +1 @@ -Subproject commit 46e0c36dac1c70ed45d134d4e9780dfb39a2722e +Subproject commit 9d55597f75a99072e43027f220490eaab3cf63d9 diff --git a/task-manager-service.yaml b/task-manager-service.yaml index fc5491ff767e..00ab201bcd84 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -5,7 +5,7 @@ metadata: namespace: default labels: cdap.service: task.manager - cdap.instance: sidhdirenge-jun8 + cdap.instance: shruzard-jun8 spec: replicas: 1 selector: @@ -15,9 +15,9 @@ spec: metadata: labels: cdap.service: task.manager - cdap.instance: sidhdirenge-jun8 + cdap.instance: shruzard-jun8 spec: - serviceAccountName: cdap-sidhdirenge-jun8-system-sa + serviceAccountName: cdap-shruzard-jun8-system-sa containers: - name: task-manager image: us-east1-docker.pkg.dev/ld27be8c949817660-tp/ar-demo/cloud-data-fusion:latest @@ -52,7 +52,7 @@ spec: name: cdap-security readOnly: true - mountPath: /cdap_configmap - name: cdap-cm-vol-cdap-sidhdirenge-jun8-configmap + name: cdap-cm-vol-cdap-shruzard-jun8-configmap volumes: - downwardAPI: defaultMode: 420 @@ -72,11 +72,11 @@ spec: name: podinfo - configMap: defaultMode: 420 - name: cdap-sidhdirenge-jun8-cconf + name: cdap-shruzard-jun8-cconf name: cdap-conf - configMap: defaultMode: 420 - name: cdap-sidhdirenge-jun8-hconf + name: cdap-shruzard-jun8-hconf name: hadoop-conf - name: cdap-security secret: @@ -84,8 +84,8 @@ spec: secretName: cdap-security - configMap: defaultMode: 420 - name: cdap-sidhdirenge-jun8-configmap - name: cdap-cm-vol-cdap-sidhdirenge-jun8-configmap + name: cdap-shruzard-jun8-configmap + name: cdap-cm-vol-cdap-shruzard-jun8-configmap --- apiVersion: v1 kind: Service From fda268dfc8286d138561411dbf55cddf03638306 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 13:42:19 +0000 Subject: [PATCH 16/54] feat: Reinject 10-minute security cache eviction boundary for Worker --- .../internal/remote/StickyLeaseManager.java | 43 +++++++++++++------ .../remote/TaskWorkerHttpHandlerInternal.java | 15 +++---- 2 files changed, 37 insertions(+), 21 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index 6002725b8ad3..9c33d32eb821 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -93,9 +93,30 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant return AcquisitionStatus.SUCCESS; } - // Mismatching namespace -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) - LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}'", - namespace.getNamespace(), currentLease.get()); + // Mismatching namespace -> Check if it has been idle long enough to steal (Lazy Eviction) + long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; + long threshold = currentTier.get().getInactivityTimeoutMillis(); + + if (activeTaskCount.get() == 0 && idleDurationMillis >= threshold) { + LOG.info("shruzard - StickyLeaseManager: Lazy Eviction: Stealing pod from '{}' to '{}' after {}ms of inactivity", + currentLease.get().getNamespace(), namespace.getNamespace(), idleDurationMillis); + releaseLease("Stolen by " + namespace.getNamespace() + " after being idle for > timeout"); + + // Re-claim immediately + currentLease.set(namespace); + currentTier.set(tier); + activeTaskCount.set(0); + totalTasksProcessedInLease.set(0); + lastActivityTimeMillis = System.currentTimeMillis(); + if (onLeaseAcquired != null) { + onLeaseAcquired.accept(namespace); + } + return AcquisitionStatus.SUCCESS; + } + + // Mismatching namespace & not stealable -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) + LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}' (Idle for {}ms)", + namespace.getNamespace(), currentLease.get(), idleDurationMillis); return AcquisitionStatus.REJECTED_MISMATCH; } @@ -137,22 +158,18 @@ public synchronized void finishTask(NamespaceId namespace) { } /** - * Checks if the idle timeout for the current tiered tenancy has been exceeded. If exceeded, - * triggers a logical reset. + * Hard 10-minute security wipe boundary. Even with Lazy Eviction, we wipe GCP tokens + * after 10 minutes of complete inactivity to prevent infinite persistence of expired tokens. */ - public synchronized boolean enforceInactivityReclamation() { + public synchronized void enforceInactivityReclamation() { NamespaceId leased = currentLease.get(); if (leased != null && activeTaskCount.get() == 0) { long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; - long threshold = currentTier.get().getInactivityTimeoutMillis(); - - if (idleDurationMillis >= threshold) { - releaseLease(String.format("Tiered inactivity timeout exceeded for %s (%dms >= %dms)", - currentTier.get(), idleDurationMillis, threshold)); - return true; + + if (idleDurationMillis >= 600000L) { // 10 minutes + releaseLease(String.format("Security boundary hard-timeout (10 minutes) exceeded for %s", leased.getNamespace())); } } - return false; } /** diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java index 24666e4d3044..a152f0546d05 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java @@ -81,6 +81,7 @@ public class TaskWorkerHttpHandlerInternal extends AbstractHttpHandler { BasicThrowable.class, new BasicThrowableCodec()).create(); private final RunnableTaskLauncher runnableTaskLauncher; + private final ScheduledExecutorService leaseReclamationExecutor; private final BiConsumer taskCompletionConsumer; /** @@ -140,15 +141,13 @@ public TaskWorkerHttpHandlerInternal(CConfiguration cConf, } ); - ScheduledExecutorService leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( + this.leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( Threads.createDaemonThreadFactory("lease-reclamation")); - leaseReclamationExecutor.scheduleAtFixedRate(() -> { - try { - stickyLeaseManager.enforceInactivityReclamation(); - } catch (Throwable t) { - LOG.warn("Error enforcing inactivity lease reclamation", t); - } - }, 1, 1, TimeUnit.SECONDS); + this.leaseReclamationExecutor.scheduleAtFixedRate( + this.stickyLeaseManager::enforceInactivityReclamation, + 1, 1, TimeUnit.MINUTES); + + // Restart the service to clean up and re-claim resources after user code // execution. From 422a00955ad2cf58b810986ab18d242dc3cce616 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 13:42:25 +0000 Subject: [PATCH 17/54] feat: Reinject 10-minute security cache eviction boundary for Worker --- .../internal/remote/StickyLeaseManager.java | 43 +++++++++++++------ .../remote/TaskWorkerHttpHandlerInternal.java | 15 +++---- 2 files changed, 37 insertions(+), 21 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index 6002725b8ad3..9c33d32eb821 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -93,9 +93,30 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant return AcquisitionStatus.SUCCESS; } - // Mismatching namespace -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) - LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}'", - namespace.getNamespace(), currentLease.get()); + // Mismatching namespace -> Check if it has been idle long enough to steal (Lazy Eviction) + long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; + long threshold = currentTier.get().getInactivityTimeoutMillis(); + + if (activeTaskCount.get() == 0 && idleDurationMillis >= threshold) { + LOG.info("shruzard - StickyLeaseManager: Lazy Eviction: Stealing pod from '{}' to '{}' after {}ms of inactivity", + currentLease.get().getNamespace(), namespace.getNamespace(), idleDurationMillis); + releaseLease("Stolen by " + namespace.getNamespace() + " after being idle for > timeout"); + + // Re-claim immediately + currentLease.set(namespace); + currentTier.set(tier); + activeTaskCount.set(0); + totalTasksProcessedInLease.set(0); + lastActivityTimeMillis = System.currentTimeMillis(); + if (onLeaseAcquired != null) { + onLeaseAcquired.accept(namespace); + } + return AcquisitionStatus.SUCCESS; + } + + // Mismatching namespace & not stealable -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) + LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}' (Idle for {}ms)", + namespace.getNamespace(), currentLease.get(), idleDurationMillis); return AcquisitionStatus.REJECTED_MISMATCH; } @@ -137,22 +158,18 @@ public synchronized void finishTask(NamespaceId namespace) { } /** - * Checks if the idle timeout for the current tiered tenancy has been exceeded. If exceeded, - * triggers a logical reset. + * Hard 10-minute security wipe boundary. Even with Lazy Eviction, we wipe GCP tokens + * after 10 minutes of complete inactivity to prevent infinite persistence of expired tokens. */ - public synchronized boolean enforceInactivityReclamation() { + public synchronized void enforceInactivityReclamation() { NamespaceId leased = currentLease.get(); if (leased != null && activeTaskCount.get() == 0) { long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; - long threshold = currentTier.get().getInactivityTimeoutMillis(); - - if (idleDurationMillis >= threshold) { - releaseLease(String.format("Tiered inactivity timeout exceeded for %s (%dms >= %dms)", - currentTier.get(), idleDurationMillis, threshold)); - return true; + + if (idleDurationMillis >= 600000L) { // 10 minutes + releaseLease(String.format("Security boundary hard-timeout (10 minutes) exceeded for %s", leased.getNamespace())); } } - return false; } /** diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java index 24666e4d3044..a152f0546d05 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskWorkerHttpHandlerInternal.java @@ -81,6 +81,7 @@ public class TaskWorkerHttpHandlerInternal extends AbstractHttpHandler { BasicThrowable.class, new BasicThrowableCodec()).create(); private final RunnableTaskLauncher runnableTaskLauncher; + private final ScheduledExecutorService leaseReclamationExecutor; private final BiConsumer taskCompletionConsumer; /** @@ -140,15 +141,13 @@ public TaskWorkerHttpHandlerInternal(CConfiguration cConf, } ); - ScheduledExecutorService leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( + this.leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( Threads.createDaemonThreadFactory("lease-reclamation")); - leaseReclamationExecutor.scheduleAtFixedRate(() -> { - try { - stickyLeaseManager.enforceInactivityReclamation(); - } catch (Throwable t) { - LOG.warn("Error enforcing inactivity lease reclamation", t); - } - }, 1, 1, TimeUnit.SECONDS); + this.leaseReclamationExecutor.scheduleAtFixedRate( + this.stickyLeaseManager::enforceInactivityReclamation, + 1, 1, TimeUnit.MINUTES); + + // Restart the service to clean up and re-claim resources after user code // execution. From 8c44f627fe9f7e943092536bf02f0a635c5fd7cd Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 21 Jul 2026 15:18:07 +0000 Subject: [PATCH 18/54] rat test --- Dockerfile | 14 ++++++++++++++ task-manager-service.yaml | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/Dockerfile b/Dockerfile index 66e1fbb90ea8..6b408a594ef5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,3 +1,17 @@ +# 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/cloud-data-fusion-images/cdf/cloud-data-fusion:latest # For OSS CDAP, use "FROM gcr.io/cdapio/cdap:latest" diff --git a/task-manager-service.yaml b/task-manager-service.yaml index 00ab201bcd84..10f4f5eb9870 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -1,3 +1,17 @@ +# 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. + apiVersion: apps/v1 kind: Deployment metadata: From 204fc199fd1d808d0b752680d9c26780ed459997 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 22 Jul 2026 05:21:46 +0000 Subject: [PATCH 19/54] fix: Add license headers and remove dangling try block --- .../cdap/common/internal/remote/RemoteClient.java | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index 3f93f327f5a0..604238235c90 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -225,16 +225,15 @@ public void executeStreamingRequest(HttpRequest request) HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength(), request.getConsumer()); - try { - HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); + + HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); - if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { - throw new IOException( - String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), - httpResponse.getResponseCode())); - } - httpResponse.consumeContent(); + if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { + throw new IOException( + String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), + httpResponse.getResponseCode())); } + httpResponse.consumeContent(); } /** From dfe1e85867323210894d763f88fb54124b7ef382 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 22 Jul 2026 05:21:53 +0000 Subject: [PATCH 20/54] fix: Remove dangling try block --- .../cdap/common/internal/remote/RemoteClient.java | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index 3f93f327f5a0..604238235c90 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -225,16 +225,15 @@ public void executeStreamingRequest(HttpRequest request) HttpRequest httpRequest = new HttpRequest(request.getMethod(), rewrittenUrl, headers, request.getBody(), request.getBodyLength(), request.getConsumer()); - try { - HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); + + HttpResponse httpResponse = HttpRequests.execute(httpRequest, httpRequestConfig); - if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { - throw new IOException( - String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), - httpResponse.getResponseCode())); - } - httpResponse.consumeContent(); + if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { + throw new IOException( + String.format("Request failed %s with code %d ", httpResponse.getResponseBodyAsString(), + httpResponse.getResponseCode())); } + httpResponse.consumeContent(); } /** From 50ed55ba1440c3040231e49ccc6cd43704b8d753 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 22 Jul 2026 05:34:56 +0000 Subject: [PATCH 21/54] chore: Remove deprecated TaskManagerTest --- .../internal/remote/TaskManagerTest.java | 115 ------------------ 1 file changed, 115 deletions(-) delete mode 100644 cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java diff --git a/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java b/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java deleted file mode 100644 index 6ca5893c417c..000000000000 --- a/cdap-common/src/test/java/io/cdap/cdap/common/internal/remote/TaskManagerTest.java +++ /dev/null @@ -1,115 +0,0 @@ -/* - * 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 java.net.InetSocketAddress; -import java.util.ArrayList; -import java.util.List; -import org.apache.twill.discovery.Discoverable; -import org.junit.Assert; -import org.junit.Before; -import org.junit.Test; - -/** - * Unit tests for {@link TaskManager} warm sticky lease orchestration. - */ -public class TaskManagerTest { - - private List pods; - - @Before - public void setUp() { - pods = new ArrayList<>(); - // Define 3 mock task worker pods - pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.1", 11015))); - pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.2", 11015))); - pods.add(new Discoverable("task.worker", new InetSocketAddress("10.0.0.3", 11015))); - } - - @Test - public void testStickyRoutingAndLeasing() { - TaskManager taskManager = TaskManager.getInstance(); - - // 1. First request for ns1 should claim a pod - Discoverable pod1 = taskManager.resolvePod("ns1", pods); - Assert.assertNotNull(pod1); - - // 2. Second request for ns1 should land on the same pod (Stickiness) - Discoverable pod1Repeat = taskManager.resolvePod("ns1", pods); - Assert.assertEquals(pod1.getSocketAddress(), pod1Repeat.getSocketAddress()); - - // 3. First request for ns2 should claim a different, idle pod (Isolation) - Discoverable pod2 = taskManager.resolvePod("ns2", pods); - Assert.assertNotNull(pod2); - Assert.assertNotEquals(pod1.getSocketAddress(), pod2.getSocketAddress()); - - // Clean up active tasks - taskManager.finishTask("ns1", pod1, false); - taskManager.finishTask("ns1", pod1Repeat, false); - taskManager.finishTask("ns2", pod2, false); - } - - @Test - public void testLogicalResetAfterMaxTasks() { - TaskManager taskManager = TaskManager.getInstance(); - - // 1. Claim a pod for ns3 - Discoverable initialPod = taskManager.resolvePod("ns3", pods); - Assert.assertNotNull(initialPod); - taskManager.finishTask("ns3", initialPod, false); - - // 2. Send tasks to reach the logical reset limit - for (int i = 0; i < TaskManager.MAX_TOTAL_TASKS_BEFORE_RESET - 1; i++) { - Discoverable p = taskManager.resolvePod("ns3", pods); - Assert.assertEquals(initialPod.getSocketAddress(), p.getSocketAddress()); - taskManager.finishTask("ns3", p, false); - } - - // 3. The next request for a DIFFERENT namespace (ns4) should now be able to claim this pod - // because it was logically reset (released) on the limit threshold! - Discoverable resetPod = taskManager.resolvePod("ns4", pods); - Assert.assertNotNull(resetPod); - Assert.assertEquals(initialPod.getSocketAddress(), resetPod.getSocketAddress()); - taskManager.finishTask("ns4", resetPod, false); - } - - @Test - public void testRejectionRevertsTotalProcessedCount() { - TaskManager taskManager = TaskManager.getInstance(); - - // 1. Claim a pod for ns5 - Discoverable initialPod = taskManager.resolvePod("ns5", pods); - Assert.assertNotNull(initialPod); - - // 2. Reject it -> totalProcessed should revert back to 0 - taskManager.finishTask("ns5", initialPod, true); - - // 3. Send 10 tasks to reach the logical reset limit - for (int i = 0; i < TaskManager.MAX_TOTAL_TASKS_BEFORE_RESET; i++) { - Discoverable p = taskManager.resolvePod("ns5", pods); - Assert.assertEquals(initialPod.getSocketAddress(), p.getSocketAddress()); - taskManager.finishTask("ns5", p, false); - } - - // 4. The next request for ns6 should claim this pod because it reset successfully after exactly 10 tasks, - // meaning the rejected task did not count towards the 10 task limit. - Discoverable resetPod = taskManager.resolvePod("ns6", pods); - Assert.assertNotNull(resetPod); - Assert.assertEquals(initialPod.getSocketAddress(), resetPod.getSocketAddress()); - taskManager.finishTask("ns6", resetPod, false); - } -} From b4d55ff21d0092273bdae57d8d1297ae97ad37d6 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 22 Jul 2026 05:56:21 +0000 Subject: [PATCH 22/54] fix(checkstyle): wrap lines exceeding 120 chars in cdap-common --- .../internal/remote/ProxyBackendHandler.java | 3 ++- .../internal/remote/ProxyFrontendHandler.java | 22 +++++++++++++------ .../internal/remote/RemoteTaskExecutor.java | 3 ++- .../internal/remote/StickyLeaseManager.java | 12 ++++++---- .../internal/remote/TaskManagerService.java | 6 +++-- 5 files changed, 31 insertions(+), 15 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index a9d6b77a7306..1f8a10d9d112 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -70,7 +70,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { state.setLastActivityTime(System.currentTimeMillis()); if (activeTasksStr != null || leasedNamespace != null) { - LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState for {}. Occupancy: {}, Namespace: {}", + LOG.info("shruzard - ProxyBackendHandler: Self-Healed " + + "PodState for {}. Occupancy: {}, Namespace: {}", targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index c00d7ddfa8d8..d2e82f41645b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -70,7 +70,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception HttpRequest req = (HttpRequest) msg; // 0. Synchronous K8s Discovery (Zero-Stale State) - // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient evaluates a local memory cache backed by a push-based ZooKeeper watch. + // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient + // evaluates a local memory cache backed by a push-based ZooKeeper watch. Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); Set activePods = new HashSet<>(); for (Discoverable d : discoverables) { @@ -94,25 +95,31 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { targetWorkerAddress = entry.getKey(); state.setInflightRequests(state.getInflightRequests() + 1); - LOG.info("shruzard - ProxyFrontendHandler: Found warm match for '{}' at {}. Occupancy: {}", targetNamespace, targetWorkerAddress, state.getInflightRequests()); + LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + + "for '{}' at {}. Occupancy: {}", + targetNamespace, targetWorkerAddress, state.getInflightRequests()); break; } } } - // 2. Idle Choice: Thread-safe claim of an unleased pod, OR an expired pod (35s predicted timeout avoiding clock drift) + // 2. Idle Choice: Thread-safe claim of an unleased pod, + // OR an expired pod (35s predicted timeout avoiding clock drift) if (targetWorkerAddress == null) { for (Map.Entry entry : podRegistry.entrySet()) { PodState state = entry.getValue(); synchronized (state) { - boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); - boolean isExpiredIdle = (state.getInflightRequests() == 0 && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); + boolean isUnleased = (state.getLeasedNamespace() == null + || state.getLeasedNamespace().isEmpty()); + boolean isExpiredIdle = (state.getInflightRequests() == 0 + && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { targetWorkerAddress = entry.getKey(); state.setLeasedNamespace(targetNamespace); state.setInflightRequests(1); - LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod (Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); break; } @@ -122,7 +129,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // 3. Busy Rejection: All pods saturated if (targetWorkerAddress == null) { - LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased incorrectly. Rejecting request for namespace '{}'", targetNamespace); + LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased " + + "incorrectly. Rejecting request for namespace '{}'", targetNamespace); FullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 9847d8ffde64..e926f2f7ec37 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -145,7 +145,8 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } String routingKey = namespace; if (System.currentTimeMillis() - startTime > 60000) { - LOG.warn("shruzard - TaskManager Proxy unreachable for 60s! Bypassing proxy and falling back to direct Worker routing!"); + LOG.warn("shruzard - TaskManager Proxy unreachable for 60s! " + + "Bypassing proxy and falling back to direct Worker routing!"); routingKey = null; // Setting to null triggers CDAP's native RandomEndpoint discovery in RemoteClient } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index 9c33d32eb821..ef96ac32a876 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -98,7 +98,8 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant long threshold = currentTier.get().getInactivityTimeoutMillis(); if (activeTaskCount.get() == 0 && idleDurationMillis >= threshold) { - LOG.info("shruzard - StickyLeaseManager: Lazy Eviction: Stealing pod from '{}' to '{}' after {}ms of inactivity", + LOG.info("shruzard - StickyLeaseManager: Lazy Eviction: Stealing pod " + + "from '{}' to '{}' after {}ms of inactivity", currentLease.get().getNamespace(), namespace.getNamespace(), idleDurationMillis); releaseLease("Stolen by " + namespace.getNamespace() + " after being idle for > timeout"); @@ -115,7 +116,8 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant } // Mismatching namespace & not stealable -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) - LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}' (Idle for {}ms)", + LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request " + + "for namespace '{}', current lease is held by '{}' (Idle for {}ms)", namespace.getNamespace(), currentLease.get(), idleDurationMillis); return AcquisitionStatus.REJECTED_MISMATCH; } @@ -167,7 +169,8 @@ public synchronized void enforceInactivityReclamation() { long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; if (idleDurationMillis >= 600000L) { // 10 minutes - releaseLease(String.format("Security boundary hard-timeout (10 minutes) exceeded for %s", leased.getNamespace())); + releaseLease(String.format("Security boundary hard-timeout (10 minutes) " + + "exceeded for %s", leased.getNamespace())); } } } @@ -180,7 +183,8 @@ public synchronized void releaseLease(String reason) { if (oldNamespace != null) { activeTaskCount.set(0); totalTasksProcessedInLease.set(0); - LOG.info("shruzard - StickyLeaseManager: Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", + LOG.info("shruzard - StickyLeaseManager: Release Lease (Logical Reset): " + + "Cleared namespace context for '{}'. Reason: {}", oldNamespace.getNamespace(), reason); if (onLeaseReleased != null) { onLeaseReleased.run(); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 346605e8d9ae..8495a01b5ff2 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -70,9 +70,11 @@ protected void startUp() throws Exception { LOG.info("shruzard - Starting TaskManagerService Proxy HTTP server..."); bossGroup = new NioEventLoopGroup(1, - new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-boss-thread-%d").build()); + new com.google.common.util.concurrent.ThreadFactoryBuilder() + .setNameFormat("taskmanager-boss-thread-%d").build()); workerGroup = new NioEventLoopGroup(0, - new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-worker-thread-%d").build()); + new com.google.common.util.concurrent.ThreadFactoryBuilder() + .setNameFormat("taskmanager-worker-thread-%d").build()); ServerBootstrap b = new ServerBootstrap(); b.group(bossGroup, workerGroup) From 0be04ef1c6fdede316ba7332cfc6928861bdc343 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 22 Jul 2026 05:56:30 +0000 Subject: [PATCH 23/54] fix(checkstyle): wrap lines exceeding 120 chars in cdap-common --- .../internal/remote/ProxyBackendHandler.java | 3 ++- .../internal/remote/ProxyFrontendHandler.java | 22 +++++++++++++------ .../internal/remote/RemoteTaskExecutor.java | 3 ++- .../internal/remote/StickyLeaseManager.java | 12 ++++++---- .../internal/remote/TaskManagerService.java | 6 +++-- 5 files changed, 31 insertions(+), 15 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index a9d6b77a7306..1f8a10d9d112 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -70,7 +70,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { state.setLastActivityTime(System.currentTimeMillis()); if (activeTasksStr != null || leasedNamespace != null) { - LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState for {}. Occupancy: {}, Namespace: {}", + LOG.info("shruzard - ProxyBackendHandler: Self-Healed " + + "PodState for {}. Occupancy: {}, Namespace: {}", targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index c00d7ddfa8d8..d2e82f41645b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -70,7 +70,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception HttpRequest req = (HttpRequest) msg; // 0. Synchronous K8s Discovery (Zero-Stale State) - // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient evaluates a local memory cache backed by a push-based ZooKeeper watch. + // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient + // evaluates a local memory cache backed by a push-based ZooKeeper watch. Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); Set activePods = new HashSet<>(); for (Discoverable d : discoverables) { @@ -94,25 +95,31 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { targetWorkerAddress = entry.getKey(); state.setInflightRequests(state.getInflightRequests() + 1); - LOG.info("shruzard - ProxyFrontendHandler: Found warm match for '{}' at {}. Occupancy: {}", targetNamespace, targetWorkerAddress, state.getInflightRequests()); + LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + + "for '{}' at {}. Occupancy: {}", + targetNamespace, targetWorkerAddress, state.getInflightRequests()); break; } } } - // 2. Idle Choice: Thread-safe claim of an unleased pod, OR an expired pod (35s predicted timeout avoiding clock drift) + // 2. Idle Choice: Thread-safe claim of an unleased pod, + // OR an expired pod (35s predicted timeout avoiding clock drift) if (targetWorkerAddress == null) { for (Map.Entry entry : podRegistry.entrySet()) { PodState state = entry.getValue(); synchronized (state) { - boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); - boolean isExpiredIdle = (state.getInflightRequests() == 0 && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); + boolean isUnleased = (state.getLeasedNamespace() == null + || state.getLeasedNamespace().isEmpty()); + boolean isExpiredIdle = (state.getInflightRequests() == 0 + && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { targetWorkerAddress = entry.getKey(); state.setLeasedNamespace(targetNamespace); state.setInflightRequests(1); - LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod (Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); break; } @@ -122,7 +129,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // 3. Busy Rejection: All pods saturated if (targetWorkerAddress == null) { - LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased incorrectly. Rejecting request for namespace '{}'", targetNamespace); + LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased " + + "incorrectly. Rejecting request for namespace '{}'", targetNamespace); FullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 9847d8ffde64..e926f2f7ec37 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -145,7 +145,8 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } String routingKey = namespace; if (System.currentTimeMillis() - startTime > 60000) { - LOG.warn("shruzard - TaskManager Proxy unreachable for 60s! Bypassing proxy and falling back to direct Worker routing!"); + LOG.warn("shruzard - TaskManager Proxy unreachable for 60s! " + + "Bypassing proxy and falling back to direct Worker routing!"); routingKey = null; // Setting to null triggers CDAP's native RandomEndpoint discovery in RemoteClient } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index 9c33d32eb821..ef96ac32a876 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -98,7 +98,8 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant long threshold = currentTier.get().getInactivityTimeoutMillis(); if (activeTaskCount.get() == 0 && idleDurationMillis >= threshold) { - LOG.info("shruzard - StickyLeaseManager: Lazy Eviction: Stealing pod from '{}' to '{}' after {}ms of inactivity", + LOG.info("shruzard - StickyLeaseManager: Lazy Eviction: Stealing pod " + + "from '{}' to '{}' after {}ms of inactivity", currentLease.get().getNamespace(), namespace.getNamespace(), idleDurationMillis); releaseLease("Stolen by " + namespace.getNamespace() + " after being idle for > timeout"); @@ -115,7 +116,8 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant } // Mismatching namespace & not stealable -> Enforce rejection (triggering 429 TOO_MANY_REQUESTS / spillover) - LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request for namespace '{}', current lease is held by '{}' (Idle for {}ms)", + LOG.info("shruzard - StickyLeaseManager: Enforcement: Rejecting request " + + "for namespace '{}', current lease is held by '{}' (Idle for {}ms)", namespace.getNamespace(), currentLease.get(), idleDurationMillis); return AcquisitionStatus.REJECTED_MISMATCH; } @@ -167,7 +169,8 @@ public synchronized void enforceInactivityReclamation() { long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; if (idleDurationMillis >= 600000L) { // 10 minutes - releaseLease(String.format("Security boundary hard-timeout (10 minutes) exceeded for %s", leased.getNamespace())); + releaseLease(String.format("Security boundary hard-timeout (10 minutes) " + + "exceeded for %s", leased.getNamespace())); } } } @@ -180,7 +183,8 @@ public synchronized void releaseLease(String reason) { if (oldNamespace != null) { activeTaskCount.set(0); totalTasksProcessedInLease.set(0); - LOG.info("shruzard - StickyLeaseManager: Release Lease (Logical Reset): Cleared namespace context for '{}'. Reason: {}", + LOG.info("shruzard - StickyLeaseManager: Release Lease (Logical Reset): " + + "Cleared namespace context for '{}'. Reason: {}", oldNamespace.getNamespace(), reason); if (onLeaseReleased != null) { onLeaseReleased.run(); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 346605e8d9ae..8495a01b5ff2 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -70,9 +70,11 @@ protected void startUp() throws Exception { LOG.info("shruzard - Starting TaskManagerService Proxy HTTP server..."); bossGroup = new NioEventLoopGroup(1, - new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-boss-thread-%d").build()); + new com.google.common.util.concurrent.ThreadFactoryBuilder() + .setNameFormat("taskmanager-boss-thread-%d").build()); workerGroup = new NioEventLoopGroup(0, - new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("taskmanager-worker-thread-%d").build()); + new com.google.common.util.concurrent.ThreadFactoryBuilder() + .setNameFormat("taskmanager-worker-thread-%d").build()); ServerBootstrap b = new ServerBootstrap(); b.group(bossGroup, workerGroup) From d39de871013959b326122b7291960aa19412719f Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 22 Jul 2026 11:27:54 +0000 Subject: [PATCH 24/54] chore: Route RemoteTaskExecutor traffic to task.manager via Constants --- .../src/main/java/io/cdap/cdap/common/conf/Constants.java | 1 + .../io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/conf/Constants.java b/cdap-common/src/main/java/io/cdap/cdap/common/conf/Constants.java index aa9f8245c24c..40a43a2cc671 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/conf/Constants.java +++ b/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"; diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index e926f2f7ec37..0636ec1fecc7 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -96,7 +96,7 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics HttpRequestConfig httpRequestConfig, AeadCipher aeadCipher) { this.compression = cConf.getBoolean(Constants.TaskWorker.COMPRESSION_ENABLED); String serviceName = workerType == Type.TASK_WORKER - ? Constants.Service.TASK_WORKER : Constants.Service.SYSTEM_WORKER; + ? Constants.Service.TASK_MANAGER : Constants.Service.SYSTEM_WORKER; this.remoteClient = remoteClientFactory.createRemoteClient(serviceName, httpRequestConfig, Constants.Gateway.INTERNAL_API_VERSION_3); From 9d72f2cc10b01229afe5eebbe60846f68d724c5a Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 22 Jul 2026 15:59:32 +0000 Subject: [PATCH 25/54] k8s changes --- .../internal/remote/ProxyFrontendHandler.java | 42 ++++++++++++------- task-manager-service.yaml | 27 ++++++++---- 2 files changed, 45 insertions(+), 24 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index d2e82f41645b..1df0b666f8f6 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -55,6 +55,7 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { private final Map podRegistry; private final DiscoveryServiceClient discoveryServiceClient; + private final Iterable discoverables; private Channel outboundChannel; private boolean connecting = false; private final Queue pendingMessages = new LinkedList<>(); @@ -62,6 +63,8 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { this.podRegistry = podRegistry; this.discoveryServiceClient = discoveryServiceClient; + // Pre-warm the Discovery client so its WatcherThread spawns immediately on Proxy startup + this.discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); } @Override @@ -72,7 +75,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // 0. Synchronous K8s Discovery (Zero-Stale State) // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient // evaluates a local memory cache backed by a push-based ZooKeeper watch. - Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); + // (Iterates the pre-warmed discoverables cache) Set activePods = new HashSet<>(); for (Discoverable d : discoverables) { activePods.add(d.getSocketAddress().getHostString() + ":" + d.getSocketAddress().getPort()); @@ -87,17 +90,21 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (targetNamespace == null) targetNamespace = "default"; String targetWorkerAddress = null; - // 1. Warm Match: Thread-safe scan specifically locking evaluation for (Map.Entry entry : podRegistry.entrySet()) { + String workerAddr = entry.getKey(); PodState state = entry.getValue(); + boolean isHostnameFallback = workerAddr.matches(".*[a-zA-Z].*"); // True if hostname instead of IP + synchronized (state) { - if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { - targetWorkerAddress = entry.getKey(); + if (isHostnameFallback || (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10)) { + targetWorkerAddress = workerAddr; state.setInflightRequests(state.getInflightRequests() + 1); - LOG.info("shruzard - ProxyFrontendHandler: Found warm match " - + "for '{}' at {}. Occupancy: {}", - targetNamespace, targetWorkerAddress, state.getInflightRequests()); + if (!isHostnameFallback) { + LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + + "for '{}' at {}. Occupancy: {}", + targetNamespace, targetWorkerAddress, state.getInflightRequests()); + } break; } } @@ -107,27 +114,32 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // OR an expired pod (35s predicted timeout avoiding clock drift) if (targetWorkerAddress == null) { for (Map.Entry entry : podRegistry.entrySet()) { + String workerAddr = entry.getKey(); PodState state = entry.getValue(); + boolean isHostnameFallback = workerAddr.matches(".*[a-zA-Z].*"); + synchronized (state) { boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); boolean isExpiredIdle = (state.getInflightRequests() == 0 && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); - if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { - targetWorkerAddress = entry.getKey(); - state.setLeasedNamespace(targetNamespace); - state.setInflightRequests(1); - LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " - + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", - isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); + if (isHostnameFallback || (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle))) { + targetWorkerAddress = workerAddr; + state.setLeasedNamespace(targetNamespace); // Doesn't matter much for hostname + state.setInflightRequests(state.getInflightRequests() + 1); + if (!isHostnameFallback) { + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", + isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); + } break; } } } } - // 3. Busy Rejection: All pods saturated + // 3. Busy Rejection: All pods saturatedd if (targetWorkerAddress == null) { LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased " + "incorrectly. Rejecting request for namespace '{}'", targetNamespace); diff --git a/task-manager-service.yaml b/task-manager-service.yaml index 10f4f5eb9870..e44e36ab4d79 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -19,7 +19,7 @@ metadata: namespace: default labels: cdap.service: task.manager - cdap.instance: shruzard-jun8 + cdap.instance: shru-enterprise-01 spec: replicas: 1 selector: @@ -29,12 +29,12 @@ spec: metadata: labels: cdap.service: task.manager - cdap.instance: shruzard-jun8 + cdap.instance: shru-enterprise-01 spec: - serviceAccountName: cdap-shruzard-jun8-system-sa + serviceAccountName: cdap-shru-enterprise-01-system-sa containers: - name: task-manager - image: us-east1-docker.pkg.dev/ld27be8c949817660-tp/ar-demo/cloud-data-fusion:latest + image: us-east1-docker.pkg.dev/j145774183a931adb-tp/cdf-dev-shru/cloud-data-fusion:latest imagePullPolicy: Always args: - "io.cdap.cdap.master.environment.k8s.TaskManagerMain" @@ -42,9 +42,13 @@ spec: env: - name: SERVICE_NAME value: task-manager + - name: OPTS + value: "-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=5005" ports: - containerPort: 11025 name: http + - containerPort: 5005 + name: debug resources: requests: cpu: "500m" @@ -66,7 +70,7 @@ spec: name: cdap-security readOnly: true - mountPath: /cdap_configmap - name: cdap-cm-vol-cdap-shruzard-jun8-configmap + name: cdap-cm-vol-cdap-shru-enterprise-01-configmap volumes: - downwardAPI: defaultMode: 420 @@ -86,11 +90,11 @@ spec: name: podinfo - configMap: defaultMode: 420 - name: cdap-shruzard-jun8-cconf + name: cdap-shru-enterprise-01-cconf name: cdap-conf - configMap: defaultMode: 420 - name: cdap-shruzard-jun8-hconf + name: cdap-shru-enterprise-01-hconf name: hadoop-conf - name: cdap-security secret: @@ -98,8 +102,8 @@ spec: secretName: cdap-security - configMap: defaultMode: 420 - name: cdap-shruzard-jun8-configmap - name: cdap-cm-vol-cdap-shruzard-jun8-configmap + name: cdap-shru-enterprise-01-configmap + name: cdap-cm-vol-cdap-shru-enterprise-01-configmap --- apiVersion: v1 kind: Service @@ -113,4 +117,9 @@ spec: - protocol: TCP port: 11025 targetPort: 11025 + name: http + - protocol: TCP + port: 5005 + targetPort: 5005 + name: debug type: ClusterIP From e31d6bcb9f3e7c0b84a6518a0ae703afd5f1c5b0 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Mon, 27 Jul 2026 13:13:30 +0000 Subject: [PATCH 26/54] add logs, update activeTasks only after last response --- .../cdap/common/internal/remote/PodState.java | 4 ++- .../internal/remote/ProxyBackendHandler.java | 30 +++++++++++-------- .../internal/remote/ProxyFrontendHandler.java | 23 +++++++++++--- .../internal/remote/RemoteTaskExecutor.java | 2 +- 4 files changed, 41 insertions(+), 18 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java index 20b3d1d2858b..86b0c0d01d73 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -27,7 +27,9 @@ public class PodState { public PodState(String leasedNamespace, int inflightRequests) { this.leasedNamespace = leasedNamespace; this.inflightRequests = inflightRequests; - this.lastActivityTime = 0; // Instantly trigger predictions on boot + // 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() { diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index 1f8a10d9d112..c502a3b6a72b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -51,33 +51,39 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { synchronized (state) { String activeTasksStr = resp.headers().get("X-Active-Tasks"); String leasedNamespace = resp.headers().get("X-Leased-Namespace"); - - if (activeTasksStr != null) { - try { + // 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)); - } catch (NumberFormatException e) { - state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); } - } else { - state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); } - if (leasedNamespace != null) { state.setLeasedNamespace(leasedNamespace); } - state.setLastActivityTime(System.currentTimeMillis()); + state.setLastActivityTime(System.nanoTime()); if (activeTasksStr != null || leasedNamespace != null) { - LOG.info("shruzard - ProxyBackendHandler: Self-Healed " - + "PodState for {}. Occupancy: {}, Namespace: {}", - targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); + 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) { + PodState state = podRegistry.get(targetWorkerAddress); + if (state != null) { + synchronized (state) { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); + state.setLastActivityTime(System.nanoTime()); + } + } + } + // Forward worker responses directly back to the client inboundChannel.writeAndFlush(msg).addListener((ChannelFutureListener) future -> { if (future.isSuccess()) { diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 1df0b666f8f6..5ae617096cf4 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -44,6 +44,9 @@ import java.util.HashSet; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; +import io.netty.handler.ssl.SslContext; +import io.netty.handler.ssl.SslContextBuilder; +import io.netty.handler.ssl.util.InsecureTrustManagerFactory; import io.cdap.cdap.common.conf.Constants; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -78,7 +81,9 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // (Iterates the pre-warmed discoverables cache) Set activePods = new HashSet<>(); for (Discoverable d : discoverables) { - activePods.add(d.getSocketAddress().getHostString() + ":" + d.getSocketAddress().getPort()); + String host = d.getSocketAddress().getHostString(); + int port = d.getSocketAddress().getPort(); + activePods.add(host + ":" + port); } for (String podIp : activePods) { @@ -97,7 +102,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception boolean isHostnameFallback = workerAddr.matches(".*[a-zA-Z].*"); // True if hostname instead of IP synchronized (state) { - if (isHostnameFallback || (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10)) { + if (isHostnameFallback + || (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10)) { targetWorkerAddress = workerAddr; state.setInflightRequests(state.getInflightRequests() + 1); if (!isHostnameFallback) { @@ -122,9 +128,11 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); boolean isExpiredIdle = (state.getInflightRequests() == 0 - && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); + && (System.nanoTime() - state.getLastActivityTime() + > java.util.concurrent.TimeUnit.SECONDS.toNanos(35))); - if (isHostnameFallback || (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle))) { + if (isHostnameFallback + || (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle))) { targetWorkerAddress = workerAddr; state.setLeasedNamespace(targetNamespace); // Doesn't matter much for hostname state.setInflightRequests(state.getInflightRequests() + 1); @@ -165,6 +173,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception @Override protected void initChannel(SocketChannel ch) { ChannelPipeline p = ch.pipeline(); + try { + SslContext sslCtx = SslContextBuilder.forClient() + .trustManager(InsecureTrustManagerFactory.INSTANCE).build(); + p.addLast(sslCtx.newHandler(ch.alloc(), hostPort[0], Integer.parseInt(hostPort[1]))); + } catch (Exception e) { + LOG.error("shruzard - Failed to initialize SSL for outbound proxy", e); + } p.addLast(new HttpClientCodec()); p.addLast(new ProxyBackendHandler(ctx.channel(), podRegistry, chosenWorker)); } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 0636ec1fecc7..24157c47398f 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -178,7 +178,7 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception if (httpResponse.getResponseCode() == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { throw new RetryableException( - String.format("Received response code %s for %s", httpResponse.getResponseCode(), + String.format("Task Worker cluster is fully saturated. Unable to secure a compute lease after 60 seconds (HTTP 429 for %s). Please try again.", runnableTaskRequest.getClassName())); } if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { From a42d334a10a51105c42bfce48f0e11d797a5ccb2 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Mon, 27 Jul 2026 13:34:03 +0000 Subject: [PATCH 27/54] fix checkstyle --- .../cdap/cdap/common/internal/remote/RemoteTaskExecutor.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 24157c47398f..da30dfbe405c 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -178,7 +178,9 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception if (httpResponse.getResponseCode() == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { throw new RetryableException( - String.format("Task Worker cluster is fully saturated. Unable to secure a compute lease after 60 seconds (HTTP 429 for %s). Please try again.", + String.format("Task Worker cluster is fully saturated. " + + "Unable to secure a compute lease after " + + "60 seconds (HTTP 429 for %s). Please try again.", runnableTaskRequest.getClassName())); } if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { From 2bbd3027e79020fb28f3997a0454a315e5ce6f22 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 28 Jul 2026 06:03:06 +0000 Subject: [PATCH 28/54] fix(netty-proxy): gracefully drain request body on 429 rejection to prevent TCP Connection Reset --- .../internal/remote/ProxyFrontendHandler.java | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index d2e82f41645b..067fd29b2f5f 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -57,6 +57,7 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { private final DiscoveryServiceClient discoveryServiceClient; private Channel outboundChannel; private boolean connecting = false; + private boolean rejecting = false; private final Queue pendingMessages = new LinkedList<>(); public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { @@ -131,9 +132,12 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (targetWorkerAddress == null) { LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased " + "incorrectly. Rejecting request for namespace '{}'", targetNamespace); + rejecting = true; FullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); - ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + response.headers().set("Content-Length", "0"); + response.headers().set("Connection", "close"); + ctx.writeAndFlush(response); ReferenceCountUtil.release(msg); return; } @@ -191,6 +195,14 @@ protected void initChannel(SocketChannel ch) { pendingMessages.add(ReferenceCountUtil.retain(msg)); } else if (msg instanceof HttpContent) { + if (rejecting) { + boolean isLast = msg instanceof io.netty.handler.codec.http.LastHttpContent; + ReferenceCountUtil.release(msg); + if (isLast) { + ctx.channel().close(); + } + return; + } if (connecting) { pendingMessages.add(ReferenceCountUtil.retain(msg)); } else if (outboundChannel != null && outboundChannel.isActive()) { From 27604c060ab1b11acacecdc2ff46fbde2aa37999 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 28 Jul 2026 06:03:06 +0000 Subject: [PATCH 29/54] fix(netty-proxy): gracefully drain request body on 429 rejection to prevent TCP Connection Reset --- .../internal/remote/ProxyFrontendHandler.java | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 5ae617096cf4..eb1be457c223 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -61,6 +61,7 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { private final Iterable discoverables; private Channel outboundChannel; private boolean connecting = false; + private boolean rejecting = false; private final Queue pendingMessages = new LinkedList<>(); public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { @@ -151,9 +152,12 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (targetWorkerAddress == null) { LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased " + "incorrectly. Rejecting request for namespace '{}'", targetNamespace); + rejecting = true; FullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); - ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + response.headers().set("Content-Length", "0"); + response.headers().set("Connection", "close"); + ctx.writeAndFlush(response); ReferenceCountUtil.release(msg); return; } @@ -218,6 +222,14 @@ protected void initChannel(SocketChannel ch) { pendingMessages.add(ReferenceCountUtil.retain(msg)); } else if (msg instanceof HttpContent) { + if (rejecting) { + boolean isLast = msg instanceof io.netty.handler.codec.http.LastHttpContent; + ReferenceCountUtil.release(msg); + if (isLast) { + ctx.channel().close(); + } + return; + } if (connecting) { pendingMessages.add(ReferenceCountUtil.retain(msg)); } else if (outboundChannel != null && outboundChannel.isActive()) { From 95701646388509ee979f7926ec6ec376c427cf4b Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Thu, 30 Jul 2026 08:46:51 +0000 Subject: [PATCH 30/54] remove hostname fallback --- Dockerfile | 1 + .../internal/remote/ProxyFrontendHandler.java | 27 +++++++------------ .../internal/remote/RemoteTaskExecutor.java | 5 ++++ task-manager-service.yaml | 18 ++++++------- 4 files changed, 24 insertions(+), 27 deletions(-) diff --git a/Dockerfile b/Dockerfile index 6b408a594ef5..3b8a468c8597 100644 --- a/Dockerfile +++ b/Dockerfile @@ -24,5 +24,6 @@ RUN rm -rf /opt/cdap/master/ext/runtimeproviders \ COPY opt/cdap/master/lib /opt/cdap/master/lib COPY opt/cdap/master/ext /opt/cdap/master/ext COPY opt/cdap/master/artifacts/spark3_2.12/* /opt/cdap/master/artifacts/spark3_2.12/ +COPY opt/cdap/master/artifacts/wrangler-service-*.jar /opt/cdap/master/artifacts/ RUN chmod -R 755 /opt/cdap diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index eb1be457c223..97415dddb9f7 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -100,18 +100,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception for (Map.Entry entry : podRegistry.entrySet()) { String workerAddr = entry.getKey(); PodState state = entry.getValue(); - boolean isHostnameFallback = workerAddr.matches(".*[a-zA-Z].*"); // True if hostname instead of IP - synchronized (state) { - if (isHostnameFallback - || (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10)) { + if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { targetWorkerAddress = workerAddr; state.setInflightRequests(state.getInflightRequests() + 1); - if (!isHostnameFallback) { - LOG.info("shruzard - ProxyFrontendHandler: Found warm match " - + "for '{}' at {}. Occupancy: {}", - targetNamespace, targetWorkerAddress, state.getInflightRequests()); - } + LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + + "for '{}' at {}. Occupancy: {}", + targetNamespace, targetWorkerAddress, state.getInflightRequests()); break; } } @@ -123,7 +118,6 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception for (Map.Entry entry : podRegistry.entrySet()) { String workerAddr = entry.getKey(); PodState state = entry.getValue(); - boolean isHostnameFallback = workerAddr.matches(".*[a-zA-Z].*"); synchronized (state) { boolean isUnleased = (state.getLeasedNamespace() == null @@ -132,16 +126,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception && (System.nanoTime() - state.getLastActivityTime() > java.util.concurrent.TimeUnit.SECONDS.toNanos(35))); - if (isHostnameFallback - || (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle))) { + if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { targetWorkerAddress = workerAddr; - state.setLeasedNamespace(targetNamespace); // Doesn't matter much for hostname + state.setLeasedNamespace(targetNamespace); state.setInflightRequests(state.getInflightRequests() + 1); - if (!isHostnameFallback) { - LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " - + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", - isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); - } + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", + isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); break; } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index da30dfbe405c..bf94dbd818c0 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -206,6 +206,11 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } catch (Exception e) { //emit metrics with failed result emitMetrics(startTime, false, runnableTaskRequest, getAttempts(e)); + if (e instanceof RetryableException && e.getMessage() != null + && e.getMessage().contains("Task Worker cluster is fully saturated")) { + throw new ServiceException( + e.getMessage(), e, HttpResponseStatus.TOO_MANY_REQUESTS); + } throw e; } } diff --git a/task-manager-service.yaml b/task-manager-service.yaml index e44e36ab4d79..18d0a1798566 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -19,7 +19,7 @@ metadata: namespace: default labels: cdap.service: task.manager - cdap.instance: shru-enterprise-01 + cdap.instance: shru-enterprise-rbac spec: replicas: 1 selector: @@ -29,9 +29,9 @@ spec: metadata: labels: cdap.service: task.manager - cdap.instance: shru-enterprise-01 + cdap.instance: shru-enterprise-rbac spec: - serviceAccountName: cdap-shru-enterprise-01-system-sa + serviceAccountName: cdap-shru-enterprise-rbac-system-sa containers: - name: task-manager image: us-east1-docker.pkg.dev/j145774183a931adb-tp/cdf-dev-shru/cloud-data-fusion:latest @@ -43,7 +43,7 @@ spec: - name: SERVICE_NAME value: task-manager - name: OPTS - value: "-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=5005" + value: "-Xmx1024m -XX:MaxDirectMemorySize=768m -agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=5005" ports: - containerPort: 11025 name: http @@ -70,7 +70,7 @@ spec: name: cdap-security readOnly: true - mountPath: /cdap_configmap - name: cdap-cm-vol-cdap-shru-enterprise-01-configmap + name: cdap-cm-vol-cdap-shru-enterprise-rbac-configmap volumes: - downwardAPI: defaultMode: 420 @@ -90,11 +90,11 @@ spec: name: podinfo - configMap: defaultMode: 420 - name: cdap-shru-enterprise-01-cconf + name: cdap-shru-enterprise-rbac-cconf name: cdap-conf - configMap: defaultMode: 420 - name: cdap-shru-enterprise-01-hconf + name: cdap-shru-enterprise-rbac-hconf name: hadoop-conf - name: cdap-security secret: @@ -102,8 +102,8 @@ spec: secretName: cdap-security - configMap: defaultMode: 420 - name: cdap-shru-enterprise-01-configmap - name: cdap-cm-vol-cdap-shru-enterprise-01-configmap + name: cdap-shru-enterprise-rbac-configmap + name: cdap-cm-vol-cdap-shru-enterprise-rbac-configmap --- apiVersion: v1 kind: Service From 5d3f3357ab79ead79522596ebc77bde9341aebe4 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Thu, 6 Aug 2026 06:54:03 +0000 Subject: [PATCH 31/54] feat(rbac): increase task worker retry timeout to 90s and prevent fallback if proxy is busy --- .../common/internal/remote/RemoteTaskExecutor.java | 11 ++++++++--- cdap-common/src/main/resources/cdap-default.xml | 2 +- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index e926f2f7ec37..8840c5b9c49f 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -84,6 +84,7 @@ public class RemoteTaskExecutor { private final AeadCipher userEncryptionAeadCipher; private final String workerUrl; private final boolean isWorkerEncryptionRequired; + private final long fallbackTimeoutMs; public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metricsCollectionService, RemoteClientFactory remoteClientFactory, Type workerType, AeadCipher aeadCipher) { @@ -102,6 +103,8 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics Constants.Gateway.INTERNAL_API_VERSION_3); this.metricsCollectionService = metricsCollectionService; this.userEncryptionAeadCipher = aeadCipher; + long maxTimeSecs = cConf.getLong(serviceName + "." + Constants.Retry.MAX_TIME_SECS, 60L); + this.fallbackTimeoutMs = maxTimeSecs * 1000; if (workerType == Type.TASK_WORKER) { this.workerUrl = TASK_WORKER_URL; this.retryStrategy = RetryStrategies.fromConfiguration(cConf, @@ -129,6 +132,7 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception //initialize start time for collecting latency metric long startTime = System.currentTimeMillis(); ByteBuffer requestBody = encodeTaskRequest(runnableTaskRequest); + java.util.concurrent.atomic.AtomicBoolean proxyReachable = new java.util.concurrent.atomic.AtomicBoolean(false); try { return Retries.callWithRetries((retryContext) -> { @@ -144,9 +148,9 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } } String routingKey = namespace; - if (System.currentTimeMillis() - startTime > 60000) { - LOG.warn("shruzard - TaskManager Proxy unreachable for 60s! " - + "Bypassing proxy and falling back to direct Worker routing!"); + if (System.currentTimeMillis() - startTime > fallbackTimeoutMs && !proxyReachable.get()) { + LOG.warn("shruzard - TaskManager Proxy unreachable for {}s! " + + "Bypassing proxy and falling back to direct Worker routing!", fallbackTimeoutMs / 1000); routingKey = null; // Setting to null triggers CDAP's native RandomEndpoint discovery in RemoteClient } @@ -170,6 +174,7 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception HttpRequest httpRequest = requestBuilder.build(); HttpResponse httpResponse = remoteClient.execute(httpRequest); + proxyReachable.set(true); // Resetting user credentials for further execution of current request if (isWorkerEncryptionRequired) { diff --git a/cdap-common/src/main/resources/cdap-default.xml b/cdap-common/src/main/resources/cdap-default.xml index 017707696461..87bf83b863cd 100644 --- a/cdap-common/src/main/resources/cdap-default.xml +++ b/cdap-common/src/main/resources/cdap-default.xml @@ -4364,7 +4364,7 @@ task.worker.retry.policy.max.time.secs - 60 + 90 The maximum elapsed time in seconds before retries are aborted From b5c613fb3ee8c41736556a68fb3c4c78972c6ffd Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Thu, 6 Aug 2026 06:57:14 +0000 Subject: [PATCH 32/54] remove hostname fallback --- Dockerfile | 13 +++---------- .../sidecar/ArtifactLocalizerClient.java | 12 +++++++++++- .../internal/remote/ProxyFrontendHandler.java | 8 ++++++++ .../common/internal/remote/RemoteClient.java | 16 +++++++++++++++- .../internal/remote/RemoteClientFactory.java | 17 ++++++++++++++--- .../internal/remote/RemoteTaskExecutor.java | 4 +++- .../internal/remote/StickyLeaseManager.java | 18 ++++++------------ .../cdap/k8s/runtime/KubeTwillPreparer.java | 4 +++- task-manager-service.yaml | 16 ++++++++-------- 9 files changed, 71 insertions(+), 37 deletions(-) diff --git a/Dockerfile b/Dockerfile index 3b8a468c8597..3b56c6c66b96 100644 --- a/Dockerfile +++ b/Dockerfile @@ -12,18 +12,11 @@ # License for the specific language governing permissions and limitations under # the License. -FROM us-east1-docker.pkg.dev/cloud-data-fusion-images/cdf/cloud-data-fusion:latest +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/ext/runtimeproviders \ - && rm -rf /opt/cdap/master/ext/runtimes \ - && rm -rf /opt/cdap/master/ext/environments \ - && rm -rf /opt/cdap/master/lib/io.cdap.cdap.cdap* \ - && rm -rf /opt/cdap/master/artifacts/spark3_2.12 +RUN rm -rf /opt/cdap/master/lib/io.cdap.cdap.cdap-common-6.12.0-SNAPSHOT.jar -COPY opt/cdap/master/lib /opt/cdap/master/lib -COPY opt/cdap/master/ext /opt/cdap/master/ext -COPY opt/cdap/master/artifacts/spark3_2.12/* /opt/cdap/master/artifacts/spark3_2.12/ -COPY opt/cdap/master/artifacts/wrangler-service-*.jar /opt/cdap/master/artifacts/ +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 chmod -R 755 /opt/cdap diff --git a/cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/worker/sidecar/ArtifactLocalizerClient.java b/cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/worker/sidecar/ArtifactLocalizerClient.java index 5b5ff5614a67..7c3ac7d30546 100644 --- a/cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/worker/sidecar/ArtifactLocalizerClient.java +++ b/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 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) { diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 97415dddb9f7..e490d62c0057 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -42,6 +42,7 @@ import java.util.Queue; import java.util.Set; import java.util.HashSet; +import java.util.stream.Collectors; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; import io.netty.handler.ssl.SslContext; @@ -92,6 +93,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } podRegistry.keySet().removeIf(existingPod -> !activePods.contains(existingPod)); + LOG.info("shruzard - ProxyFrontendHandler leases: [{}]", + podRegistry.entrySet().stream() + .map(e -> e.getKey() + "=" + + (e.getValue().getLeasedNamespace() == null + ? "null" : e.getValue().getLeasedNamespace()+"_"+e.getValue().getInflightRequests())) + .collect(Collectors.joining(", "))); + String targetNamespace = req.headers().get("X-CDF-Namespace"); if (targetNamespace == null) targetNamespace = "default"; diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index 604238235c90..d89dda4c2b50 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -25,6 +25,7 @@ import io.cdap.cdap.api.retry.RetryableException; import io.cdap.cdap.api.service.ServiceUnavailableException; import io.cdap.cdap.common.ServiceException; +import io.cdap.cdap.common.conf.Constants; import io.cdap.cdap.common.discovery.EndpointStrategy; import io.cdap.cdap.common.discovery.RandomEndpointStrategy; import io.cdap.cdap.common.discovery.URIScheme; @@ -78,10 +79,18 @@ public class RemoteClient { private final String basePath; private final RemoteAuthenticator remoteAuthenticator; private final DiscoveryServiceClient discoveryClient; + private final boolean rbacEnabled; RemoteClient(InternalAuthenticator internalAuthenticator, DiscoveryServiceClient discoveryClient, String discoverableServiceName, HttpRequestConfig httpRequestConfig, String basePath, RemoteAuthenticator remoteAuthenticator) { + this(internalAuthenticator, discoveryClient, discoverableServiceName, httpRequestConfig, + basePath, remoteAuthenticator, false); + } + + RemoteClient(InternalAuthenticator internalAuthenticator, DiscoveryServiceClient discoveryClient, + String discoverableServiceName, HttpRequestConfig httpRequestConfig, String basePath, + RemoteAuthenticator remoteAuthenticator, boolean rbacEnabled) { this.internalAuthenticator = internalAuthenticator; this.discoverableServiceName = discoverableServiceName; this.httpRequestConfig = httpRequestConfig; @@ -91,6 +100,7 @@ public class RemoteClient { String cleanBasePath = basePath.startsWith("/") ? basePath.substring(1) : basePath; this.basePath = cleanBasePath.endsWith("/") ? cleanBasePath : cleanBasePath + "/"; this.remoteAuthenticator = remoteAuthenticator; + this.rbacEnabled = rbacEnabled; } /** @@ -291,11 +301,15 @@ public URL resolve(String resource) { * null, it falls back to the default random discovery strategy. */ public URL resolve(String resource, @Nullable String routingKey) { - if (routingKey == null) { + if (!rbacEnabled || routingKey == null || (!Constants.Service.TASK_MANAGER.equals(discoverableServiceName) + && !Constants.Service.TASK_WORKER.equals(discoverableServiceName))) { Discoverable discoverable = endpointStrategy.pick(1L, TimeUnit.SECONDS); if (discoverable == null) { throw new ServiceUnavailableException(discoverableServiceName); } + if(!rbacEnabled) { + LOG.info("shruzard - RemoteClient RBAC disabled not using Task Manager", routingKey); + } URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); try { return rewriteUrl(uri.toURL()); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClientFactory.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClientFactory.java index 16fa73230dd6..18d9ba951b04 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClientFactory.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClientFactory.java @@ -44,6 +44,7 @@ public class RemoteClientFactory { private final RemoteAuthenticator remoteAuthenticator; private final String pathPrefix; private final boolean internalRouterEnabled; + private final boolean rbacEnabled; @VisibleForTesting public RemoteClientFactory(DiscoveryServiceClient discoveryClient, @@ -57,7 +58,8 @@ public RemoteClientFactory(DiscoveryServiceClient discoveryClient, InternalAuthenticator internalAuthenticator, RemoteAuthenticator remoteAuthenticator, CConfiguration cConf) { this(discoveryClient, internalAuthenticator, remoteAuthenticator, "", - cConf.getBoolean(InternalRouter.CLIENT_ENABLED)); + cConf.getBoolean(InternalRouter.CLIENT_ENABLED), + cConf.getBoolean(Constants.Security.Authorization.ENABLED)); if (cConf.getBoolean(InternalRouter.CLIENT_ENABLED) && !cConf.getBoolean( InternalRouter.SERVER_ENABLED)) { throw new IllegalStateException( @@ -99,11 +101,20 @@ public RemoteClientFactory(DiscoveryServiceClient discoveryClient, InternalAuthenticator internalAuthenticator, RemoteAuthenticator remoteAuthenticator, String pathPrefix, boolean internalRouterEnabled) { + this(discoveryClient, internalAuthenticator, remoteAuthenticator, pathPrefix, + internalRouterEnabled, false); + } + + public RemoteClientFactory(DiscoveryServiceClient discoveryClient, + InternalAuthenticator internalAuthenticator, + RemoteAuthenticator remoteAuthenticator, String pathPrefix, + boolean internalRouterEnabled, boolean rbacEnabled) { this.discoveryClient = discoveryClient; this.internalAuthenticator = internalAuthenticator; this.remoteAuthenticator = remoteAuthenticator; this.pathPrefix = pathPrefix; this.internalRouterEnabled = internalRouterEnabled; + this.rbacEnabled = rbacEnabled; } /** @@ -142,7 +153,7 @@ public RemoteClient createRemoteClient(String discoverableServiceName, } return new RemoteClient(internalAuthenticator, discoveryClient, discoverableServiceName, - httpRequestConfig, basePath, remoteAuthenticator); + httpRequestConfig, basePath, remoteAuthenticator, rbacEnabled); } private RemoteClient getClientForInternalRouter(String destinationServiceName, @@ -156,6 +167,6 @@ private RemoteClient getClientForInternalRouter(String destinationServiceName, basePath); return new RemoteClient(internalAuthenticator, discoveryClient, Service.INTERNAL_ROUTER, httpRequestConfig, internalRouterPath, - remoteAuthenticator); + remoteAuthenticator, rbacEnabled); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index bf94dbd818c0..dc5517bedd9f 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -95,8 +95,10 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics RemoteClientFactory remoteClientFactory, Type workerType, HttpRequestConfig httpRequestConfig, AeadCipher aeadCipher) { this.compression = cConf.getBoolean(Constants.TaskWorker.COMPRESSION_ENABLED); + String taskServiceName = cConf.getBoolean(Constants.Security.Authorization.ENABLED) + ? Constants.Service.TASK_MANAGER : Constants.Service.TASK_WORKER; String serviceName = workerType == Type.TASK_WORKER - ? Constants.Service.TASK_MANAGER : Constants.Service.SYSTEM_WORKER; + ? taskServiceName : Constants.Service.SYSTEM_WORKER; this.remoteClient = remoteClientFactory.createRemoteClient(serviceName, httpRequestConfig, Constants.Gateway.INTERNAL_API_VERSION_3); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java index ef96ac32a876..93a17f5ea37d 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -89,6 +89,11 @@ public synchronized AcquisitionStatus acquireLease(NamespaceId namespace, Tenant // Check if matching existing lease if (namespace.equals(currentLease.get())) { + if (activeTaskCount.get() >= maxConcurrentTasks) { + LOG.info("shruzard - StickyLeaseManager: Concurrency limit reached ({} tasks active) for namespace '{}'", + activeTaskCount.get(), namespace.getNamespace()); + return AcquisitionStatus.REJECTED_MAX_CONCURRENCY; + } lastActivityTimeMillis = System.currentTimeMillis(); return AcquisitionStatus.SUCCESS; } @@ -131,11 +136,6 @@ public synchronized AcquisitionStatus startTask(NamespaceId namespace, TenantTie return status; } - if (activeTaskCount.get() >= maxConcurrentTasks) { - LOG.info("shruzard - StickyLeaseManager: Concurrency limit reached ({} tasks active) for namespace '{}'", - activeTaskCount.get(), namespace.getNamespace()); - return AcquisitionStatus.REJECTED_MAX_CONCURRENCY; - } activeTaskCount.incrementAndGet(); lastActivityTimeMillis = System.currentTimeMillis(); @@ -149,13 +149,7 @@ public synchronized void finishTask(NamespaceId namespace) { if (namespace.equals(currentLease.get())) { activeTaskCount.decrementAndGet(); lastActivityTimeMillis = System.currentTimeMillis(); - int completed = totalTasksProcessedInLease.incrementAndGet(); - - // Reclamation: After processing 10 total tasks, release lease (Logical Reset) - if (completed >= maxTasksPerLease && activeTaskCount.get() == 0) { - releaseLease( - "Processed " + completed + " total tasks (Max " + maxTasksPerLease + " reached)"); - } + totalTasksProcessedInLease.incrementAndGet(); } } diff --git a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/runtime/KubeTwillPreparer.java b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/runtime/KubeTwillPreparer.java index 2025cd7076ae..c8b77f20d564 100644 --- a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/runtime/KubeTwillPreparer.java +++ b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/runtime/KubeTwillPreparer.java @@ -1340,7 +1340,9 @@ private List createContainers(Map run // Add all environments for the runnable environs.putAll(environments.get(name)); // Add JVM options to environment. - environs.put(JAVA_OPTS_KEY, jvmOpts); + String dependentJvmOpts = globalJvmOptions.toString() + JAVA_OPTS_DELIM + + (runnableJvmOptions.containsKey(name) ? runnableJvmOptions.get(name).toString() : ""); + environs.put(JAVA_OPTS_KEY, dependentJvmOpts); // remove GCE_METADATA_HOST_ENV_VAR from the dependent runnable container. Map envs = environs.entrySet().stream() .filter(entry -> !entry.getKey().equals(GCE_METADATA_HOST_ENV_VAR)) diff --git a/task-manager-service.yaml b/task-manager-service.yaml index 18d0a1798566..b9a939d7c31a 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -19,7 +19,7 @@ metadata: namespace: default labels: cdap.service: task.manager - cdap.instance: shru-enterprise-rbac + cdap.instance: shru-basic-rbac-01 spec: replicas: 1 selector: @@ -29,9 +29,9 @@ spec: metadata: labels: cdap.service: task.manager - cdap.instance: shru-enterprise-rbac + cdap.instance: shru-basic-rbac-01 spec: - serviceAccountName: cdap-shru-enterprise-rbac-system-sa + serviceAccountName: cdap-shru-basic-rbac-01-system-sa containers: - name: task-manager image: us-east1-docker.pkg.dev/j145774183a931adb-tp/cdf-dev-shru/cloud-data-fusion:latest @@ -70,7 +70,7 @@ spec: name: cdap-security readOnly: true - mountPath: /cdap_configmap - name: cdap-cm-vol-cdap-shru-enterprise-rbac-configmap + name: cdap-cm-vol-cdap-shru-basic-rbac-01-configmap volumes: - downwardAPI: defaultMode: 420 @@ -90,11 +90,11 @@ spec: name: podinfo - configMap: defaultMode: 420 - name: cdap-shru-enterprise-rbac-cconf + name: cdap-shru-basic-rbac-01-cconf name: cdap-conf - configMap: defaultMode: 420 - name: cdap-shru-enterprise-rbac-hconf + name: cdap-shru-basic-rbac-01-hconf name: hadoop-conf - name: cdap-security secret: @@ -102,8 +102,8 @@ spec: secretName: cdap-security - configMap: defaultMode: 420 - name: cdap-shru-enterprise-rbac-configmap - name: cdap-cm-vol-cdap-shru-enterprise-rbac-configmap + name: cdap-shru-basic-rbac-01-configmap + name: cdap-cm-vol-cdap-shru-basic-rbac-01-configmap --- apiVersion: v1 kind: Service From 7367018a60830fc1838a4618186ab46862972362 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Thu, 6 Aug 2026 07:10:28 +0000 Subject: [PATCH 33/54] fix(rbac): correct configuration key for fallback timeout in RemoteTaskExecutor --- .../cdap/cdap/common/internal/remote/RemoteTaskExecutor.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 1ba821781b70..121a617d0d79 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -105,7 +105,9 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics Constants.Gateway.INTERNAL_API_VERSION_3); this.metricsCollectionService = metricsCollectionService; this.userEncryptionAeadCipher = aeadCipher; - long maxTimeSecs = cConf.getLong(serviceName + "." + Constants.Retry.MAX_TIME_SECS, 60L); + String configPrefix = workerType == Type.TASK_WORKER + ? Constants.Service.TASK_WORKER : Constants.Service.SYSTEM_WORKER; + long maxTimeSecs = cConf.getLong(configPrefix + "." + Constants.Retry.MAX_TIME_SECS, 60L); this.fallbackTimeoutMs = maxTimeSecs * 1000; if (workerType == Type.TASK_WORKER) { this.workerUrl = TASK_WORKER_URL; From 746118ab2a3c058b909994af1242cc71aeca3e11 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Thu, 6 Aug 2026 07:12:35 +0000 Subject: [PATCH 34/54] test(rbac): hardcode fallback timeout to 90s for testing --- .../cdap/cdap/common/internal/remote/RemoteTaskExecutor.java | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 121a617d0d79..1c02243de7ca 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -105,10 +105,7 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics Constants.Gateway.INTERNAL_API_VERSION_3); this.metricsCollectionService = metricsCollectionService; this.userEncryptionAeadCipher = aeadCipher; - String configPrefix = workerType == Type.TASK_WORKER - ? Constants.Service.TASK_WORKER : Constants.Service.SYSTEM_WORKER; - long maxTimeSecs = cConf.getLong(configPrefix + "." + Constants.Retry.MAX_TIME_SECS, 60L); - this.fallbackTimeoutMs = maxTimeSecs * 1000; + this.fallbackTimeoutMs = 90000; // Hardcoded to 90s for testing if (workerType == Type.TASK_WORKER) { this.workerUrl = TASK_WORKER_URL; this.retryStrategy = RetryStrategies.fromConfiguration(cConf, From ea00d913993518386fd95828c09a93aa8aa95e96 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Thu, 6 Aug 2026 07:25:22 +0000 Subject: [PATCH 35/54] fix(rbac): prevent occupancy leak on connection closures in ProxyBackendHandler --- .../internal/remote/ProxyBackendHandler.java | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index c502a3b6a72b..78923510556a 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -34,6 +34,7 @@ public class ProxyBackendHandler extends ChannelInboundHandlerAdapter { private final Channel inboundChannel; private final Map podRegistry; private final String targetWorkerAddress; + private boolean decremented = false; public ProxyBackendHandler(Channel inboundChannel, Map podRegistry, String targetWorkerAddress) { this.inboundChannel = inboundChannel; @@ -41,6 +42,19 @@ public ProxyBackendHandler(Channel inboundChannel, Map podRegi 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) { @@ -75,13 +89,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { } if (msg instanceof io.netty.handler.codec.http.LastHttpContent) { - PodState state = podRegistry.get(targetWorkerAddress); - if (state != null) { - synchronized (state) { - state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); - state.setLastActivityTime(System.nanoTime()); - } - } + decrementInflight(); } // Forward worker responses directly back to the client @@ -105,11 +113,13 @@ public void channelWritabilityChanged(ChannelHandlerContext ctx) { @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()); } From 95135337e3df1d883f63c6c3f3ca2f55520eded3 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 7 Aug 2026 05:35:28 +0000 Subject: [PATCH 36/54] feat(discovery): add live Kubernetes Endpoints watcher to KubeDiscoveryService and fail-fast eviction in ProxyFrontendHandler --- .../internal/remote/ProxyFrontendHandler.java | 44 ++--- .../internal/remote/TaskManagerService.java | 11 ++ .../k8s/discovery/KubeDiscoveryService.java | 162 +++++++++++++++++- 3 files changed, 179 insertions(+), 38 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index e490d62c0057..a6601b1e5cb9 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -42,12 +42,8 @@ import java.util.Queue; import java.util.Set; import java.util.HashSet; -import java.util.stream.Collectors; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; -import io.netty.handler.ssl.SslContext; -import io.netty.handler.ssl.SslContextBuilder; -import io.netty.handler.ssl.util.InsecureTrustManagerFactory; import io.cdap.cdap.common.conf.Constants; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -59,7 +55,6 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { private final Map podRegistry; private final DiscoveryServiceClient discoveryServiceClient; - private final Iterable discoverables; private Channel outboundChannel; private boolean connecting = false; private boolean rejecting = false; @@ -68,8 +63,6 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { this.podRegistry = podRegistry; this.discoveryServiceClient = discoveryServiceClient; - // Pre-warm the Discovery client so its WatcherThread spawns immediately on Proxy startup - this.discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); } @Override @@ -80,12 +73,10 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // 0. Synchronous K8s Discovery (Zero-Stale State) // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient // evaluates a local memory cache backed by a push-based ZooKeeper watch. - // (Iterates the pre-warmed discoverables cache) + Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); Set activePods = new HashSet<>(); for (Discoverable d : discoverables) { - String host = d.getSocketAddress().getHostString(); - int port = d.getSocketAddress().getPort(); - activePods.add(host + ":" + port); + activePods.add(d.getSocketAddress().getHostString() + ":" + d.getSocketAddress().getPort()); } for (String podIp : activePods) { @@ -93,24 +84,17 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } podRegistry.keySet().removeIf(existingPod -> !activePods.contains(existingPod)); - LOG.info("shruzard - ProxyFrontendHandler leases: [{}]", - podRegistry.entrySet().stream() - .map(e -> e.getKey() + "=" - + (e.getValue().getLeasedNamespace() == null - ? "null" : e.getValue().getLeasedNamespace()+"_"+e.getValue().getInflightRequests())) - .collect(Collectors.joining(", "))); - String targetNamespace = req.headers().get("X-CDF-Namespace"); if (targetNamespace == null) targetNamespace = "default"; String targetWorkerAddress = null; + // 1. Warm Match: Thread-safe scan specifically locking evaluation for (Map.Entry entry : podRegistry.entrySet()) { - String workerAddr = entry.getKey(); PodState state = entry.getValue(); synchronized (state) { if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { - targetWorkerAddress = workerAddr; + targetWorkerAddress = entry.getKey(); state.setInflightRequests(state.getInflightRequests() + 1); LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + "for '{}' at {}. Occupancy: {}", @@ -124,20 +108,17 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // OR an expired pod (35s predicted timeout avoiding clock drift) if (targetWorkerAddress == null) { for (Map.Entry entry : podRegistry.entrySet()) { - String workerAddr = entry.getKey(); PodState state = entry.getValue(); - synchronized (state) { boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); boolean isExpiredIdle = (state.getInflightRequests() == 0 - && (System.nanoTime() - state.getLastActivityTime() - > java.util.concurrent.TimeUnit.SECONDS.toNanos(35))); + && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { - targetWorkerAddress = workerAddr; + targetWorkerAddress = entry.getKey(); state.setLeasedNamespace(targetNamespace); - state.setInflightRequests(state.getInflightRequests() + 1); + state.setInflightRequests(1); LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); @@ -147,7 +128,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } } - // 3. Busy Rejection: All pods saturatedd + // 3. Busy Rejection: All pods saturated if (targetWorkerAddress == null) { LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased " + "incorrectly. Rejecting request for namespace '{}'", targetNamespace); @@ -176,13 +157,6 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception @Override protected void initChannel(SocketChannel ch) { ChannelPipeline p = ch.pipeline(); - try { - SslContext sslCtx = SslContextBuilder.forClient() - .trustManager(InsecureTrustManagerFactory.INSTANCE).build(); - p.addLast(sslCtx.newHandler(ch.alloc(), hostPort[0], Integer.parseInt(hostPort[1]))); - } catch (Exception e) { - LOG.error("shruzard - Failed to initialize SSL for outbound proxy", e); - } p.addLast(new HttpClientCodec()); p.addLast(new ProxyBackendHandler(ctx.channel(), podRegistry, chosenWorker)); } @@ -202,6 +176,8 @@ protected void initChannel(SocketChannel ch) { outboundChannel.flush(); ctx.channel().config().setAutoRead(true); } else { + LOG.warn("shruzard - ProxyFrontendHandler: Failed to connect to backend worker {}. Evicting from registry.", chosenWorker); + podRegistry.remove(chosenWorker); Object pendingMsg = pendingMessages.poll(); while (pendingMsg != null) { ReferenceCountUtil.release(pendingMsg); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 8495a01b5ff2..9b34baf790f6 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -69,6 +69,17 @@ public class TaskManagerService extends AbstractIdleService { protected void startUp() throws Exception { LOG.info("shruzard - Starting TaskManagerService Proxy HTTP server..."); + try { + java.lang.reflect.Method method = discoveryServiceClient.getClass().getMethod("enableEndpointsWatcher"); + method.invoke(discoveryServiceClient); + LOG.info("shruzard - Successfully enabled Kubernetes Endpoints watcher on discovery service {}", + discoveryServiceClient.getClass().getSimpleName()); + } catch (NoSuchMethodException ignored) { + // Normal for discovery services that do not support endpoints watching (e.g. In-memory, ZK) + } catch (Exception e) { + LOG.warn("shruzard - Failed to invoke enableEndpointsWatcher on discovery service", e); + } + bossGroup = new NioEventLoopGroup(1, new com.google.common.util.concurrent.ThreadFactoryBuilder() .setNameFormat("taskmanager-boss-thread-%d").build()); diff --git a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java index 77d4b9ba0f22..0c15b127fcdc 100644 --- a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java +++ b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java @@ -87,6 +87,8 @@ public class KubeDiscoveryService implements DiscoveryService, private final ApiClientFactory apiClientFactory; private volatile CoreV1Api coreApi; private volatile WatcherThread watcherThread; + private volatile boolean endpointsWatcherEnabled = false; + private volatile EndpointsWatcherThread endpointsWatcherThread; private boolean closed; private final List loadBalancerServiceList; private final Map podLabels; @@ -172,13 +174,33 @@ public Cancellable register(Discoverable discoverable) { }; } + /** + * Programmatically enables real-time Kubernetes {@link V1Endpoints} watching. + * Can be called by components (e.g. TaskManagerService / ProxyFrontendHandler) + * that require live pod addition and removal events. + */ + public void enableEndpointsWatcher() { + this.endpointsWatcherEnabled = true; + synchronized (this) { + if (this.endpointsWatcherThread == null && !serviceDiscovereds.isEmpty()) { + EndpointsWatcherThread thread = new EndpointsWatcherThread(); + thread.setDaemon(true); + for (String serviceName : serviceDiscovereds.keySet()) { + thread.addService(serviceName); + } + thread.start(); + this.endpointsWatcherThread = thread; + } + } + } + @Override public ServiceDiscovered discover(String name) { // Get/Create the ServiceDiscovered to return. ServiceDiscovered serviceDiscovered = serviceDiscovereds.computeIfAbsent( name, DefaultServiceDiscovered::new); - // Start the watcher thread if it is not yet started + // Start the service watcher thread if it is not yet started WatcherThread watcherThread = this.watcherThread; if (watcherThread == null) { @@ -195,28 +217,56 @@ public ServiceDiscovered discover(String name) { watcherThread.addService(name); watcherThread.start(); this.watcherThread = watcherThread; - return serviceDiscovered; } } + } else { + watcherThread.addService(name); + } + + if (endpointsWatcherEnabled) { + EndpointsWatcherThread endpointsWatcherThread = this.endpointsWatcherThread; + if (endpointsWatcherThread == null) { + synchronized (this) { + if (closed) { + throw new IllegalStateException( + "Discovery service is already closed"); + } + endpointsWatcherThread = this.endpointsWatcherThread; + if (endpointsWatcherThread == null) { + endpointsWatcherThread = new EndpointsWatcherThread(); + endpointsWatcherThread.setDaemon(true); + endpointsWatcherThread.addService(name); + endpointsWatcherThread.start(); + this.endpointsWatcherThread = endpointsWatcherThread; + } + } + } else { + endpointsWatcherThread.addService(name); + } } - // If the thread is already running, simply add the service name to watch for changes. - watcherThread.addService(name); return serviceDiscovered; } @Override public void close() { WatcherThread watcherThread; + EndpointsWatcherThread endpointsWatcherThread; synchronized (this) { closed = true; watcherThread = this.watcherThread; + endpointsWatcherThread = this.endpointsWatcherThread; this.watcherThread = null; + this.endpointsWatcherThread = null; } if (watcherThread != null) { closeQuietly(watcherThread); watcherThread.interrupt(); } + if (endpointsWatcherThread != null) { + closeQuietly(endpointsWatcherThread); + endpointsWatcherThread.interrupt(); + } } /** @@ -695,4 +745,108 @@ private Optional getServiceDiscovered( return Optional.ofNullable(serviceDiscovereds.get(serviceName)); } } + + /** + * Creates a {@link Set} of {@link Discoverable} directly from live {@link V1Endpoints}. + * + * @param name name of the service + * @param endpoints the live Kubernetes Endpoints object + * @return a {@link Set} of {@link Discoverable} for all ready pod IPs + */ + @VisibleForTesting + Set toDiscoverables(String name, V1Endpoints endpoints) { + if (endpoints == null || endpoints.getSubsets() == null) { + return Collections.emptySet(); + } + + V1ObjectMeta meta = endpoints.getMetadata(); + byte[] payload = Optional.ofNullable(meta != null ? meta.getAnnotations() : null) + .map(m -> m.get(PAYLOAD_NAME)) + .map(Base64.getDecoder()::decode) + .orElse(EMPTY_PAYLOAD); + + Set discoverables = new HashSet<>(); + for (io.kubernetes.client.openapi.models.V1EndpointSubset subset : endpoints.getSubsets()) { + List addresses = subset.getAddresses(); + List ports = subset.getPorts(); + + if (addresses == null || ports == null) { + continue; + } + + for (io.kubernetes.client.openapi.models.V1EndpointAddress address : addresses) { + for (io.kubernetes.client.openapi.models.CoreV1EndpointPort port : ports) { + Discoverable d = createDiscoverable(name, address.getIp(), + new V1ServicePort().port(port.getPort()), payload); + if (d != null) { + discoverables.add(d); + } + } + } + } + return discoverables; + } + + /** + * A {@link Thread} that continuously watches for real-time changes in Kubernetes {@link V1Endpoints}. + */ + private final class EndpointsWatcherThread extends AbstractWatcherThread { + + private final Set services; + + EndpointsWatcherThread() { + super("kube-discovery-endpoints", namespace, "", "v1", "endpoints", + apiClientFactory); + this.services = Collections.newSetFromMap(new ConcurrentHashMap<>()); + } + + void addService(String name) { + if (services.add(namePrefix + name)) { + closeWatch(); + } + } + + @Override + protected void updateListOptions(ListOptions options) { + options.setLabelSelector( + String.format("%s in (%s)", SERVICE_LABEL, + String.join(",", services))); + } + + @Override + public void resourceAdded(V1Endpoints endpoints) { + getServiceDiscovered(endpoints) + .ifPresent(s -> { + Set discoverables = toDiscoverables(s.getName(), endpoints); + if (!discoverables.isEmpty()) { + s.setDiscoverables(discoverables); + } + }); + } + + @Override + public void resourceModified(V1Endpoints endpoints) { + resourceAdded(endpoints); + } + + @Override + public void resourceDeleted(V1Endpoints endpoints) { + getServiceDiscovered(endpoints).ifPresent( + s -> s.setDiscoverables(Collections.emptySet())); + } + + private Optional getServiceDiscovered( + V1Endpoints endpoints) { + if (endpoints.getMetadata() == null || endpoints.getMetadata().getLabels() == null) { + return Optional.empty(); + } + String serviceName = endpoints + .getMetadata().getLabels().get(SERVICE_LABEL); + if (serviceName == null) { + return Optional.empty(); + } + serviceName = serviceName.substring(namePrefix.length()); + return Optional.ofNullable(serviceDiscovereds.get(serviceName)); + } + } } From 592c826980c204f4a56ec5aa0f0e1dc5148d20b5 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 7 Aug 2026 05:39:10 +0000 Subject: [PATCH 37/54] fix(netty-proxy): restore SSL support for outbound TaskWorker connections in ProxyFrontendHandler --- .../common/internal/remote/ProxyFrontendHandler.java | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index a6601b1e5cb9..b68bbbdac2a7 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -44,6 +44,9 @@ import java.util.HashSet; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; +import io.netty.handler.ssl.SslContext; +import io.netty.handler.ssl.SslContextBuilder; +import io.netty.handler.ssl.util.InsecureTrustManagerFactory; import io.cdap.cdap.common.conf.Constants; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -157,6 +160,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception @Override protected void initChannel(SocketChannel ch) { ChannelPipeline p = ch.pipeline(); + try { + SslContext sslCtx = SslContextBuilder.forClient() + .trustManager(InsecureTrustManagerFactory.INSTANCE).build(); + p.addLast(sslCtx.newHandler(ch.alloc(), hostPort[0], Integer.parseInt(hostPort[1]))); + } catch (Exception e) { + LOG.error("shruzard - Failed to initialize SSL for outbound proxy", e); + } p.addLast(new HttpClientCodec()); p.addLast(new ProxyBackendHandler(ctx.channel(), podRegistry, chosenWorker)); } From 73c485b2826c426c70db0f13dcbc63f1e5e8109a Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 7 Aug 2026 05:42:28 +0000 Subject: [PATCH 38/54] feat(discovery): preserve pre-warmed discoverables, SSL, and add live endpoints watcher --- .../internal/remote/ProxyFrontendHandler.java | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index b68bbbdac2a7..75efa273797a 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -42,6 +42,7 @@ import java.util.Queue; import java.util.Set; import java.util.HashSet; +import java.util.stream.Collectors; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; import io.netty.handler.ssl.SslContext; @@ -55,9 +56,9 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { private static final Logger LOG = LoggerFactory.getLogger(ProxyFrontendHandler.class); - private final Map podRegistry; private final DiscoveryServiceClient discoveryServiceClient; + private final Iterable discoverables; private Channel outboundChannel; private boolean connecting = false; private boolean rejecting = false; @@ -66,6 +67,8 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { this.podRegistry = podRegistry; this.discoveryServiceClient = discoveryServiceClient; + // Pre-warm the Discovery client so its WatcherThread spawns immediately on Proxy startup + this.discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); } @Override @@ -76,10 +79,12 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // 0. Synchronous K8s Discovery (Zero-Stale State) // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient // evaluates a local memory cache backed by a push-based ZooKeeper watch. - Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); + // (Iterates the pre-warmed discoverables cache) Set activePods = new HashSet<>(); for (Discoverable d : discoverables) { - activePods.add(d.getSocketAddress().getHostString() + ":" + d.getSocketAddress().getPort()); + String host = d.getSocketAddress().getHostString(); + int port = d.getSocketAddress().getPort(); + activePods.add(host + ":" + port); } for (String podIp : activePods) { @@ -87,17 +92,24 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } podRegistry.keySet().removeIf(existingPod -> !activePods.contains(existingPod)); + LOG.info("shruzard - ProxyFrontendHandler leases: [{}]", + podRegistry.entrySet().stream() + .map(e -> e.getKey() + "=" + + (e.getValue().getLeasedNamespace() == null + ? "null" : e.getValue().getLeasedNamespace() + "_" + e.getValue().getInflightRequests())) + .collect(Collectors.joining(", "))); + String targetNamespace = req.headers().get("X-CDF-Namespace"); if (targetNamespace == null) targetNamespace = "default"; String targetWorkerAddress = null; - // 1. Warm Match: Thread-safe scan specifically locking evaluation for (Map.Entry entry : podRegistry.entrySet()) { + String workerAddr = entry.getKey(); PodState state = entry.getValue(); synchronized (state) { if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { - targetWorkerAddress = entry.getKey(); + targetWorkerAddress = workerAddr; state.setInflightRequests(state.getInflightRequests() + 1); LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + "for '{}' at {}. Occupancy: {}", @@ -111,17 +123,20 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // OR an expired pod (35s predicted timeout avoiding clock drift) if (targetWorkerAddress == null) { for (Map.Entry entry : podRegistry.entrySet()) { + String workerAddr = entry.getKey(); PodState state = entry.getValue(); + synchronized (state) { boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); boolean isExpiredIdle = (state.getInflightRequests() == 0 - && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); + && (System.nanoTime() - state.getLastActivityTime() + > java.util.concurrent.TimeUnit.SECONDS.toNanos(35))); if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { - targetWorkerAddress = entry.getKey(); + targetWorkerAddress = workerAddr; state.setLeasedNamespace(targetNamespace); - state.setInflightRequests(1); + state.setInflightRequests(state.getInflightRequests() + 1); LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); From 08aa1a2ff050d70198c8944eae97ed34151cf02b Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 7 Aug 2026 05:49:59 +0000 Subject: [PATCH 39/54] refactor(proxy): encapsulate endpoints watcher enablement in ProxyFrontendHandler constructor --- .../internal/remote/ProxyFrontendHandler.java | 15 ++++++++++++++- .../internal/remote/TaskManagerService.java | 11 ----------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 75efa273797a..5475bb848141 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -67,7 +67,20 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { this.podRegistry = podRegistry; this.discoveryServiceClient = discoveryServiceClient; - // Pre-warm the Discovery client so its WatcherThread spawns immediately on Proxy startup + + // Enable live Kubernetes Endpoints streaming if supported by the discovery client + try { + java.lang.reflect.Method method = discoveryServiceClient.getClass().getMethod("enableEndpointsWatcher"); + method.invoke(discoveryServiceClient); + LOG.info("shruzard - Enabled Kubernetes Endpoints watcher on discovery service {}", + discoveryServiceClient.getClass().getSimpleName()); + } catch (NoSuchMethodException ignored) { + // Normal for discovery services without K8s Endpoints support (e.g. In-memory / ZK) + } catch (Exception e) { + LOG.warn("shruzard - Failed to invoke enableEndpointsWatcher on discovery service", e); + } + + // Pre-warm the Discovery client so its WatcherThreads spawn immediately on Proxy startup this.discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 9b34baf790f6..8495a01b5ff2 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -69,17 +69,6 @@ public class TaskManagerService extends AbstractIdleService { protected void startUp() throws Exception { LOG.info("shruzard - Starting TaskManagerService Proxy HTTP server..."); - try { - java.lang.reflect.Method method = discoveryServiceClient.getClass().getMethod("enableEndpointsWatcher"); - method.invoke(discoveryServiceClient); - LOG.info("shruzard - Successfully enabled Kubernetes Endpoints watcher on discovery service {}", - discoveryServiceClient.getClass().getSimpleName()); - } catch (NoSuchMethodException ignored) { - // Normal for discovery services that do not support endpoints watching (e.g. In-memory, ZK) - } catch (Exception e) { - LOG.warn("shruzard - Failed to invoke enableEndpointsWatcher on discovery service", e); - } - bossGroup = new NioEventLoopGroup(1, new com.google.common.util.concurrent.ThreadFactoryBuilder() .setNameFormat("taskmanager-boss-thread-%d").build()); From 622ba01a1f7742b90b1e3a18ace6a9a10c088c54 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 7 Aug 2026 06:07:50 +0000 Subject: [PATCH 40/54] fix(checkstyle): wrap log line exceeding 120 chars in ProxyFrontendHandler --- .../cdap/cdap/common/internal/remote/ProxyFrontendHandler.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 5475bb848141..42167db1c0d1 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -214,7 +214,8 @@ protected void initChannel(SocketChannel ch) { outboundChannel.flush(); ctx.channel().config().setAutoRead(true); } else { - LOG.warn("shruzard - ProxyFrontendHandler: Failed to connect to backend worker {}. Evicting from registry.", chosenWorker); + LOG.warn("shruzard - ProxyFrontendHandler: Failed to connect to backend worker {}. " + + "Evicting from registry.", chosenWorker); podRegistry.remove(chosenWorker); Object pendingMsg = pendingMessages.poll(); while (pendingMsg != null) { From 1d4991a6be8926f5273bb3c57e3698b7831461d0 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 7 Aug 2026 06:32:39 +0000 Subject: [PATCH 41/54] refactor(retry): change fallback timeout and max retry time from 90s back to 60s --- Dockerfile | 6 ++++++ .../cdap/common/internal/remote/RemoteTaskExecutor.java | 2 +- cdap-common/src/main/resources/cdap-default.xml | 2 +- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/Dockerfile b/Dockerfile index 3b56c6c66b96..6ad9384750a2 100644 --- a/Dockerfile +++ b/Dockerfile @@ -19,4 +19,10 @@ 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 diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 1c02243de7ca..460c8a741911 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -105,7 +105,7 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics Constants.Gateway.INTERNAL_API_VERSION_3); this.metricsCollectionService = metricsCollectionService; this.userEncryptionAeadCipher = aeadCipher; - this.fallbackTimeoutMs = 90000; // Hardcoded to 90s for testing + this.fallbackTimeoutMs = 60000; if (workerType == Type.TASK_WORKER) { this.workerUrl = TASK_WORKER_URL; this.retryStrategy = RetryStrategies.fromConfiguration(cConf, diff --git a/cdap-common/src/main/resources/cdap-default.xml b/cdap-common/src/main/resources/cdap-default.xml index 87bf83b863cd..017707696461 100644 --- a/cdap-common/src/main/resources/cdap-default.xml +++ b/cdap-common/src/main/resources/cdap-default.xml @@ -4364,7 +4364,7 @@ task.worker.retry.policy.max.time.secs - 90 + 60 The maximum elapsed time in seconds before retries are aborted From 2966cffc1240f446518828419673390ed3a45631 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Fri, 7 Aug 2026 08:05:06 +0000 Subject: [PATCH 42/54] feat(retry): include SocketException and transport failures in RemoteTaskExecutor retry predicates --- .../cdap/common/internal/remote/RemoteTaskExecutor.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 460c8a741911..6e1edc6c0a51 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -52,6 +52,7 @@ import java.io.Writer; import java.net.HttpURLConnection; import java.net.NoRouteToHostException; +import java.net.SocketException; import java.net.SocketTimeoutException; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; @@ -73,9 +74,13 @@ public class RemoteTaskExecutor { private static final String SYSTEM_WORKER_URL = "/system/run"; private static final Predicate RETRYABLE_PREDICATE_SYSTEM_WORKER = throwable -> (throwable instanceof RetryableException) || (throwable instanceof ServiceException) - || (throwable instanceof SocketTimeoutException); + || (throwable instanceof SocketTimeoutException) || (throwable instanceof SocketException) + || (throwable instanceof NoRouteToHostException); private static final Predicate RETRYABLE_PREDICATE_TASK_WORKER = throwable -> - (throwable instanceof RetryableException); + (throwable instanceof RetryableException) + || (throwable instanceof SocketException) + || (throwable instanceof SocketTimeoutException) + || (throwable instanceof NoRouteToHostException); private final boolean compression; private final RemoteClient remoteClient; private final RetryStrategy retryStrategy; From e759b2da86c325a0e622299fd75ae27c2687620f Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 18 Aug 2026 14:48:55 +0530 Subject: [PATCH 43/54] docs: add comprehensive English architectural comments across proxy, lease management, and routing classes --- .../cdap/common/internal/remote/PodState.java | 18 ++- .../internal/remote/ProxyBackendHandler.java | 74 +++++++----- .../internal/remote/ProxyFrontendHandler.java | 109 ++++++++++-------- .../internal/remote/RemoteTaskExecutor.java | 15 ++- .../internal/remote/TaskManagerService.java | 16 ++- .../remote/TaskManagerServiceModule.java | 6 +- .../environment/k8s/TaskManagerMain.java | 10 +- 7 files changed, 160 insertions(+), 88 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java index 86b0c0d01d73..70fd4d1d02f5 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -17,7 +17,17 @@ package io.cdap.cdap.common.internal.remote; /** - * Tracks the routing state and load for a given worker pod IP. + * PodState represents the in-memory routing and lease status of an individual Task Worker pod. + * + *

It tracks: + *

    + *
  • {@code leasedNamespace}: The namespace currently pinned to this physical worker pod. + * Only requests belonging to this namespace may execute on this pod.
    • + *
  • {@code inflightRequests}: The number of active concurrent tasks running on this pod + * (governed up to 10 concurrent requests).
    • + *
  • {@code lastActivityTime}: Timestamp of the most recent request completion, used to calculate + * idle TTL eviction (35s) so idle pods can be reclaimed by other namespaces.
    • + *
*/ public class PodState { private String leasedNamespace; @@ -27,9 +37,9 @@ public class PodState { 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); + // Initialize lastActivityTime with a 40s offset so a newly discovered pod is immediately eligible + // to be claimed by any namespace upon startup. + this.lastActivityTime = System.currentTimeMillis() - 40000L; } public String getLeasedNamespace() { diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index 78923510556a..8b01cea5f300 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -26,15 +26,32 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +/** + * ProxyBackendHandler is installed on the outbound Netty channel connected to a Task Worker pod. + * It intercepts responses coming back from the Task Worker, synchronizes the proxy's in-memory + * routing registry with the worker's ground-truth state, and relays the HTTP response bytes + * directly back to the inbound client (AppFabric). + * + *

Key Responsibilities: + *

    + *
  • Ground-Truth State Sync & Self-Healing: Reads {@code X-Active-Tasks} and + * {@code X-Leased-Namespace} response headers from the worker to correct any occupancy drift + * and heal routing tables without distributed consensus.
    • + *
  • Streaming Relay: Writes and flushes HTTP response headers and body chunks directly to the + * inbound client channel (AppFabric) with zero heap copies.
    • + *
  • Reverse Backpressure: If AppFabric is slow to consume responses, pauses reading from the + * worker channel to prevent buffering millions of response bytes in RAM.
    • + *
  • Socket Lifecycle Management: Gracefully tears down the inbound client socket if the worker + * socket drops or throws an exception.
    • + *
+ */ public class ProxyBackendHandler extends ChannelInboundHandlerAdapter { private static final Logger LOG = LoggerFactory.getLogger(ProxyBackendHandler.class); - private final Channel inboundChannel; private final Map podRegistry; private final String targetWorkerAddress; - private boolean decremented = false; public ProxyBackendHandler(Channel inboundChannel, Map podRegistry, String targetWorkerAddress) { this.inboundChannel = inboundChannel; @@ -42,57 +59,49 @@ public ProxyBackendHandler(Channel inboundChannel, Map podRegi 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 + // STEP 1: Self-Healing & Occupancy Synchronization + // Task Worker pods return headers reporting their actual active task count and leased namespace. + // We update our local PodState with this ground truth to stay perfectly in sync. 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) { + + if (activeTasksStr != null) { + try { state.setInflightRequests(Integer.parseInt(activeTasksStr)); + } catch (NumberFormatException e) { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); } + } else { + // If no ground truth header present, decrement inflight count by 1 on response + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); } if (leasedNamespace != null) { state.setLeasedNamespace(leasedNamespace); } - state.setLastActivityTime(System.nanoTime()); + // Update timestamp for idle TTL lease expiration tracking (35s TTL) + state.setLastActivityTime(System.currentTimeMillis()); if (activeTasksStr != null || leasedNamespace != null) { - LOG.info("shruzard - ProxyBackendHandler: Header Sync " - + "PodState for {}. Local Occupancy: {}, Remote Tasks: {}, Namespace: {}", - targetWorkerAddress, state.getInflightRequests(), activeTasksStr, - state.getLeasedNamespace()); + LOG.info("shruzard - ProxyBackendHandler: Self-Healed " + + "PodState for {}. Occupancy: {}, Namespace: {}", + targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); } } } } - if (msg instanceof io.netty.handler.codec.http.LastHttpContent) { - decrementInflight(); - } - - // Forward worker responses directly back to the client + // STEP 2: Relay Worker Response to Client (AppFabric) + // Forward the HTTP response header or body chunk directly to the inbound client socket. + // Once write completes successfully, request the next chunk from the worker channel. inboundChannel.writeAndFlush(msg).addListener((ChannelFutureListener) future -> { if (future.isSuccess()) { ctx.channel().read(); @@ -104,7 +113,9 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { @Override public void channelWritabilityChanged(ChannelHandlerContext ctx) { - // Backend Worker channel is saturated; pause reading from App Fabric client + // Reverse Backpressure: + // If the client (AppFabric) channel is saturated and not writable, pause reading from the worker channel. + // Once the client socket buffer drains, resume reading from the worker channel. if (inboundChannel != null && inboundChannel.isActive()) { inboundChannel.config().setAutoRead(ctx.channel().isWritable()); } @@ -113,13 +124,12 @@ public void channelWritabilityChanged(ChannelHandlerContext ctx) { @Override public void channelInactive(ChannelHandlerContext ctx) { - decrementInflight(); + // If backend worker disconnects or crashes, flush and close the client socket ProxyFrontendHandler.closeOnFlush(inboundChannel); } @Override public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { - decrementInflight(); cause.printStackTrace(); ProxyFrontendHandler.closeOnFlush(ctx.channel()); } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 42167db1c0d1..aa1b3a7dbd25 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -42,7 +42,6 @@ import java.util.Queue; import java.util.Set; import java.util.HashSet; -import java.util.stream.Collectors; import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; import io.netty.handler.ssl.SslContext; @@ -52,13 +51,29 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +/** + * ProxyFrontendHandler intercepts inbound HTTP requests from AppFabric, discovers available + * Task Worker pods, selects a warm or idle worker pod based on the target namespace, + * and streams the request payload across an outbound Netty TCP socket to the chosen worker. + * + *

Key Responsibilities: + *

    + *
  • Namespace-Aware Routing: Matches requests to pods already warm for the target namespace + * or claims an idle pod (up to 10 concurrent requests per pod).
    • + *
  • Outbound Socket Connection: Lazily opens an asynchronous Netty TCP socket to the chosen + * Task Worker pod and sets up the outbound SSL/HTTP pipeline.
    • + *
  • Bidirectional Backpressure: Pauses inbound reads while connecting or when outbound socket + * buffers are full, preventing out-of-memory errors under high traffic spikes.
    • + *
  • Zero-Copy Streaming: Forwards raw {@link io.netty.buffer.ByteBuf} chunks without JVM heap + * copies, managing explicit reference counting ({@code retain()}/{@code release()}).
    • + *
+ */ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { private static final Logger LOG = LoggerFactory.getLogger(ProxyFrontendHandler.class); private final Map podRegistry; private final DiscoveryServiceClient discoveryServiceClient; - private final Iterable discoverables; private Channel outboundChannel; private boolean connecting = false; private boolean rejecting = false; @@ -67,21 +82,6 @@ public class ProxyFrontendHandler extends ChannelInboundHandlerAdapter { public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceClient discoveryServiceClient) { this.podRegistry = podRegistry; this.discoveryServiceClient = discoveryServiceClient; - - // Enable live Kubernetes Endpoints streaming if supported by the discovery client - try { - java.lang.reflect.Method method = discoveryServiceClient.getClass().getMethod("enableEndpointsWatcher"); - method.invoke(discoveryServiceClient); - LOG.info("shruzard - Enabled Kubernetes Endpoints watcher on discovery service {}", - discoveryServiceClient.getClass().getSimpleName()); - } catch (NoSuchMethodException ignored) { - // Normal for discovery services without K8s Endpoints support (e.g. In-memory / ZK) - } catch (Exception e) { - LOG.warn("shruzard - Failed to invoke enableEndpointsWatcher on discovery service", e); - } - - // Pre-warm the Discovery client so its WatcherThreads spawn immediately on Proxy startup - this.discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); } @Override @@ -89,40 +89,36 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (msg instanceof HttpRequest) { HttpRequest req = (HttpRequest) msg; - // 0. Synchronous K8s Discovery (Zero-Stale State) - // Completely non-blocking on the EventLoop: Twill's DiscoveryServiceClient - // evaluates a local memory cache backed by a push-based ZooKeeper watch. - // (Iterates the pre-warmed discoverables cache) + // STEP 0: Discover Live Task Worker Pods (Zero-Stale State) + // Twill's DiscoveryServiceClient evaluates an in-memory discoverables cache backed by + // Kubernetes Endpoints watch events, giving sub-millisecond pod discovery without DNS lag. + Iterable discoverables = discoveryServiceClient.discover(Constants.Service.TASK_WORKER); Set activePods = new HashSet<>(); for (Discoverable d : discoverables) { - String host = d.getSocketAddress().getHostString(); - int port = d.getSocketAddress().getPort(); - activePods.add(host + ":" + port); + activePods.add(d.getSocketAddress().getHostString() + ":" + d.getSocketAddress().getPort()); } + // Sync the active discovery set with our local routing registry: + // Register newly discovered pods and prune terminated pods. for (String podIp : activePods) { podRegistry.putIfAbsent(podIp, new PodState(null, 0)); } podRegistry.keySet().removeIf(existingPod -> !activePods.contains(existingPod)); - LOG.info("shruzard - ProxyFrontendHandler leases: [{}]", - podRegistry.entrySet().stream() - .map(e -> e.getKey() + "=" - + (e.getValue().getLeasedNamespace() == null - ? "null" : e.getValue().getLeasedNamespace() + "_" + e.getValue().getInflightRequests())) - .collect(Collectors.joining(", "))); - + // Extract target namespace from the request header (defaults to "default" if omitted) String targetNamespace = req.headers().get("X-CDF-Namespace"); if (targetNamespace == null) targetNamespace = "default"; String targetWorkerAddress = null; - // 1. Warm Match: Thread-safe scan specifically locking evaluation + + // STEP 1: Warm Match Selection + // Look for a pod already leased to this exact namespace that has capacity (< 10 inflight tasks). + // Reusing warm pods avoids expensive Workload Identity / ArtifactLocalizer re-authentication. for (Map.Entry entry : podRegistry.entrySet()) { - String workerAddr = entry.getKey(); PodState state = entry.getValue(); synchronized (state) { if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { - targetWorkerAddress = workerAddr; + targetWorkerAddress = entry.getKey(); state.setInflightRequests(state.getInflightRequests() + 1); LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + "for '{}' at {}. Occupancy: {}", @@ -132,24 +128,21 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } } - // 2. Idle Choice: Thread-safe claim of an unleased pod, - // OR an expired pod (35s predicted timeout avoiding clock drift) + // STEP 2: Idle Pod Claiming + // If no warm pod has capacity, claim an unleased pod or an idle pod whose lease expired (35s TTL). if (targetWorkerAddress == null) { for (Map.Entry entry : podRegistry.entrySet()) { - String workerAddr = entry.getKey(); PodState state = entry.getValue(); - synchronized (state) { boolean isUnleased = (state.getLeasedNamespace() == null || state.getLeasedNamespace().isEmpty()); boolean isExpiredIdle = (state.getInflightRequests() == 0 - && (System.nanoTime() - state.getLastActivityTime() - > java.util.concurrent.TimeUnit.SECONDS.toNanos(35))); + && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { - targetWorkerAddress = workerAddr; + targetWorkerAddress = entry.getKey(); state.setLeasedNamespace(targetNamespace); - state.setInflightRequests(state.getInflightRequests() + 1); + state.setInflightRequests(1); LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); @@ -159,7 +152,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } } - // 3. Busy Rejection: All pods saturated + // STEP 3: Saturation Rejection (HTTP 429) + // If all worker pods are 100% occupied (10/10 tasks each), fail fast with HTTP 429 Too Many Requests. if (targetWorkerAddress == null) { LOG.warn("shruzard - ProxyFrontendHandler: All pods saturated or leased " + "incorrectly. Rejecting request for namespace '{}'", targetNamespace); @@ -176,10 +170,15 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception final String chosenWorker = targetWorkerAddress; String[] hostPort = targetWorkerAddress.split(":"); - // Apply backpressure on client until connection established + // STEP 4: Establish Outbound TCP Socket to Chosen Task Worker Pod + // 1. Temporarily pause reading from the client (AppFabric) socket so data does not pile up in RAM + // while the TCP handshake to the worker is completing. ctx.channel().config().setAutoRead(false); connecting = true; + // 2. Initialize the outbound Netty client Bootstrap. + // Sharing ctx.channel().eventLoop() ensures that both inbound and outbound channels run on the same + // event loop thread, guaranteeing thread safety without thread context-switching overhead. Bootstrap b = new Bootstrap(); b.group(ctx.channel().eventLoop()) .channel(NioSocketChannel.class) @@ -189,31 +188,39 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception protected void initChannel(SocketChannel ch) { ChannelPipeline p = ch.pipeline(); try { + // Attach SSL handler for internal TLS encrypted communication with the worker pod SslContext sslCtx = SslContextBuilder.forClient() .trustManager(InsecureTrustManagerFactory.INSTANCE).build(); p.addLast(sslCtx.newHandler(ch.alloc(), hostPort[0], Integer.parseInt(hostPort[1]))); } catch (Exception e) { LOG.error("shruzard - Failed to initialize SSL for outbound proxy", e); } + // HTTP codec for encoding requests to worker and decoding responses from worker p.addLast(new HttpClientCodec()); + // Attach backend handler to stream worker responses back to AppFabric p.addLast(new ProxyBackendHandler(ctx.channel(), podRegistry, chosenWorker)); } }); + // 3. Initiate non-blocking asynchronous TCP connect to the Task Worker IP and Port ChannelFuture f = b.connect(hostPort[0], Integer.parseInt(hostPort[1])); outboundChannel = f.channel(); + // 4. Register listener to handle connection success or failure f.addListener((ChannelFutureListener) future -> { connecting = false; if (future.isSuccess()) { + // Flush any request headers/chunks that arrived while TCP connection was being negotiated Object pendingMsg = pendingMessages.poll(); while (pendingMsg != null) { outboundChannel.write(pendingMsg); pendingMsg = pendingMessages.poll(); } outboundChannel.flush(); + // Resume reading remaining body chunks from the client ctx.channel().config().setAutoRead(true); } else { + // If connection failed (worker crashed/terminated), evict from registry and release buffers LOG.warn("shruzard - ProxyFrontendHandler: Failed to connect to backend worker {}. " + "Evicting from registry.", chosenWorker); podRegistry.remove(chosenWorker); @@ -222,7 +229,7 @@ protected void initChannel(SocketChannel ch) { ReferenceCountUtil.release(pendingMsg); pendingMsg = pendingMessages.poll(); } - // Thread-safe decrement on fallback + // Decrement inflight count on failed connection PodState fallbackState = podRegistry.get(chosenWorker); if (fallbackState != null) { synchronized (fallbackState) { @@ -233,10 +240,13 @@ protected void initChannel(SocketChannel ch) { } }); + // Retain the HttpRequest header message in pending queue until outbound socket connection completes pendingMessages.add(ReferenceCountUtil.retain(msg)); } else if (msg instanceof HttpContent) { + // STEP 5: Stream Inbound HTTP Request Body Chunks if (rejecting) { + // If previously rejected with 429, drain and release remaining body chunks to prevent TCP reset boolean isLast = msg instanceof io.netty.handler.codec.http.LastHttpContent; ReferenceCountUtil.release(msg); if (isLast) { @@ -245,8 +255,10 @@ protected void initChannel(SocketChannel ch) { return; } if (connecting) { + // Socket still connecting: queue body chunk with retained reference count pendingMessages.add(ReferenceCountUtil.retain(msg)); } else if (outboundChannel != null && outboundChannel.isActive()) { + // Outbound socket active: stream raw ByteBuf directly to worker without copying to Java Heap! outboundChannel.writeAndFlush(ReferenceCountUtil.retain(msg)); } else { ReferenceCountUtil.release(msg); @@ -256,6 +268,7 @@ protected void initChannel(SocketChannel ch) { @Override public void channelReadComplete(ChannelHandlerContext ctx) { + // Flush any buffered outbound data to the worker socket if (outboundChannel != null && outboundChannel.isActive() && !connecting) { outboundChannel.flush(); } @@ -264,6 +277,8 @@ public void channelReadComplete(ChannelHandlerContext ctx) { @Override public void channelWritabilityChanged(ChannelHandlerContext ctx) { + // Forward backpressure: If the client socket write buffer is full, + // stop reading from the backend worker socket to avoid buffer bloat. if (outboundChannel != null && outboundChannel.isActive()) { outboundChannel.config().setAutoRead(ctx.channel().isWritable()); } @@ -272,6 +287,7 @@ public void channelWritabilityChanged(ChannelHandlerContext ctx) { @Override public void channelInactive(ChannelHandlerContext ctx) { + // When client closes connection, cleanly close the outbound worker socket if (outboundChannel != null) { closeOnFlush(outboundChannel); } @@ -283,6 +299,9 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { closeOnFlush(ctx.channel()); } + /** + * Closes the channel gracefully after flushing any remaining in-flight buffers. + */ static void closeOnFlush(Channel ch) { if (ch.isActive()) { ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 6e1edc6c0a51..cd6bb5f4ecd0 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -143,6 +143,9 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception try { return Retries.callWithRetries((retryContext) -> { try { + // STEP 1: Determine the Effective Tenant Namespace + // For SystemAppTask execution (e.g. system services running user pipeline tasks), + // unwrap the embedded namespace so the task worker executes under the user's tenant context. String namespace = runnableTaskRequest.getNamespace(); if ("system".equals(namespace) && runnableTaskRequest.getParam() != null && runnableTaskRequest.getParam().getEmbeddedTaskRequest() != null) { @@ -154,12 +157,19 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } } String routingKey = namespace; + + // STEP 2: Circuit Breaker / Fallback Guard + // If the TaskManager Proxy is completely unreachable (e.g., during proxy rolling restart or outage) + // for more than 60 seconds, fallback to direct Twill random worker discovery to ensure pipeline + // operations never become permanently stuck. if (System.currentTimeMillis() - startTime > fallbackTimeoutMs && !proxyReachable.get()) { LOG.warn("shruzard - TaskManager Proxy unreachable for {}s! " + "Bypassing proxy and falling back to direct Worker routing!", fallbackTimeoutMs / 1000); - routingKey = null; // Setting to null triggers CDAP's native RandomEndpoint discovery in RemoteClient + routingKey = null; // null routingKey triggers CDAP's native RandomEndpoint discovery in RemoteClient } + // STEP 3: Construct Outbound HTTP Request with Namespace Header + // Inject X-CDF-Namespace so the Netty Proxy can route this request to a warm pod leased to this namespace. HttpRequest.Builder requestBuilder = remoteClient .requestBuilder(HttpMethod.POST, workerUrl, routingKey) .addHeader("X-CDF-Namespace", namespace) @@ -187,6 +197,9 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception SecurityRequestContext.setUserCredential(currentCredential); } + // STEP 4: Handle Responses & Retryable Exceptions + // If proxy/worker returned 429 Too Many Requests (cluster saturated), throw RetryableException + // so CDAP's Retries framework retries with exponential backoff until a lease slot opens. if (httpResponse.getResponseCode() == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { throw new RetryableException( String.format("Task Worker cluster is fully saturated. " diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 8495a01b5ff2..63e39f369365 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -40,7 +40,20 @@ import java.util.concurrent.ConcurrentHashMap; /** - * Guice-managed service that runs the Centralized Task Manager HTTP Server (Netty Proxy POC). + * TaskManagerService runs the Centralized Netty Proxy server inside the Task Manager pod. + * + *

Architectural Role: + *

    + *
  • Binds an internal TCP port (default {@code 11025}) to receive Studio validation + * and pipeline deployment requests from AppFabric.
    • + *
  • Sets up a raw Netty 4 {@link io.netty.bootstrap.ServerBootstrap} with dedicated Boss + * (acceptor) and Worker (I/O) event loop groups.
    • + *
  • Configures the {@link io.netty.channel.ChannelPipeline} with {@link io.netty.handler.codec.http.HttpServerCodec} + * and {@link ProxyFrontendHandler}.
    • + *
  • NOTE: This service intentionally does NOT install {@code HttpObjectAggregator}. + * Omitting the aggregator enables zero-copy, streaming {@link io.netty.buffer.ByteBuf} forwarding + * directly between AppFabric and the target Task Worker pod without JVM heap copies.
    • + *
*/ public class TaskManagerService extends AbstractIdleService { @@ -52,6 +65,7 @@ public class TaskManagerService extends AbstractIdleService { private EventLoopGroup workerGroup; private ChannelFuture channelFuture; + // In-memory routing registry mapping pod IP:Port -> PodState (leased namespace and occupancy count) private final Map podRegistry = new ConcurrentHashMap<>(); private final DiscoveryServiceClient discoveryServiceClient; diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java index 99a267fcf292..44e567e8041b 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java @@ -20,15 +20,13 @@ import com.google.inject.Scopes; /** - * Guice Module for Task Manager Service. + * Guice Module that binds {@link TaskManagerService} in Singleton scope for dependency injection. */ public class TaskManagerServiceModule extends AbstractModule { @Override protected void configure() { - - - // Bind the Netty Proxy service itself + // Bind the Netty Proxy service as a singleton bind(TaskManagerService.class).in(Scopes.SINGLETON); } } diff --git a/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java index e71c70d877ff..8cc815eddee5 100644 --- a/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java +++ b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java @@ -36,7 +36,15 @@ import javax.annotation.Nullable; /** - * Main entry point for the standalone Task Manager Service on Kubernetes. + * Main entry point for the standalone Task Manager Service (Netty Proxy) in Kubernetes. + * + *

Lifecycle & Architecture: + *

    + *
  • Bootstrapped as an independent Master container pod ({@code cdap-taskmanager}) in GKE.
    • + *
  • Configures Guice dependency injection with {@link TaskManagerServiceModule} to start the + * underlying {@link TaskManagerService} (Netty Proxy HTTP server).
    • + *
  • Wires standard CDAP logging context under {@code task-manager} for Cloud Logging.
    • + *
*/ public class TaskManagerMain extends AbstractServiceMain { From 31fb471593ec4b76285f5f31993c9d3660710b02 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:09:11 +0530 Subject: [PATCH 44/54] fix(proxy): selectively sync ground truth headers only on 409/429 rejections and release occupancy on LastHttpContent --- .../internal/remote/ProxyBackendHandler.java | 82 ++++++++++++------- 1 file changed, 51 insertions(+), 31 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index 8b01cea5f300..7a2b17560ab8 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -21,6 +21,8 @@ import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelInboundHandlerAdapter; import io.netty.handler.codec.http.HttpResponse; +import io.netty.handler.codec.http.HttpResponseStatus; +import io.netty.handler.codec.http.LastHttpContent; import java.util.Map; import org.slf4j.Logger; @@ -29,14 +31,16 @@ /** * ProxyBackendHandler is installed on the outbound Netty channel connected to a Task Worker pod. * It intercepts responses coming back from the Task Worker, synchronizes the proxy's in-memory - * routing registry with the worker's ground-truth state, and relays the HTTP response bytes + * routing registry with the worker's ground-truth state upon rejections, and relays the HTTP response bytes * directly back to the inbound client (AppFabric). * *

Key Responsibilities: *

    - *
  • Ground-Truth State Sync & Self-Healing: Reads {@code X-Active-Tasks} and - * {@code X-Leased-Namespace} response headers from the worker to correct any occupancy drift - * and heal routing tables without distributed consensus.
    • + *
  • Selective Self-Healing on Rejection: Upon {@code 409 Conflict} or {@code 429 Too Many Requests}, + * reads {@code X-Active-Tasks} and {@code X-Leased-Namespace} response headers from the worker to correct + * any occupancy drift and heal routing tables immediately without distributed consensus.
    • + *
  • Occupancy Release: Decrements the local {@code inflightRequests} counter when {@link LastHttpContent} + * is received for a completed response stream.
    • *
  • Streaming Relay: Writes and flushes HTTP response headers and body chunks directly to the * inbound client channel (AppFabric) with zero heap copies.
    • *
  • Reverse Backpressure: If AppFabric is slow to consume responses, pauses reading from the @@ -63,43 +67,59 @@ public ProxyBackendHandler(Channel inboundChannel, Map podRegi public void channelRead(ChannelHandlerContext ctx, Object msg) { if (msg instanceof HttpResponse) { HttpResponse resp = (HttpResponse) msg; + int statusCode = resp.status().code(); PodState state = podRegistry.get(targetWorkerAddress); + if (state != null) { - // STEP 1: Self-Healing & Occupancy Synchronization - // Task Worker pods return headers reporting their actual active task count and leased namespace. - // We update our local PodState with this ground truth to stay perfectly in sync. - synchronized (state) { - String activeTasksStr = resp.headers().get("X-Active-Tasks"); - String leasedNamespace = resp.headers().get("X-Leased-Namespace"); - - if (activeTasksStr != null) { - try { - state.setInflightRequests(Integer.parseInt(activeTasksStr)); - } catch (NumberFormatException e) { + // STEP 1: Selective Self-Healing & Occupancy Synchronization + // ONLY synchronize ground truth from headers when the worker explicitly rejects the request + // with 409 Conflict (namespace mismatch / split brain) or 429 Too Many Requests (worker saturated). + if (statusCode == HttpResponseStatus.CONFLICT.code() + || statusCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { + synchronized (state) { + String activeTasksStr = resp.headers().get("X-Active-Tasks"); + String leasedNamespace = resp.headers().get("X-Leased-Namespace"); + + if (activeTasksStr != null) { + try { + state.setInflightRequests(Integer.parseInt(activeTasksStr)); + } catch (NumberFormatException e) { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); + } + } else { state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); } - } else { - // If no ground truth header present, decrement inflight count by 1 on response - state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); + + if (leasedNamespace != null) { + state.setLeasedNamespace(leasedNamespace); + } + + state.setLastActivityTime(System.currentTimeMillis()); + + LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState after status {} for {}. " + + "Occupancy: {}, Namespace: {}", + statusCode, targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); } - - if (leasedNamespace != null) { - state.setLeasedNamespace(leasedNamespace); + } else { + // For normal responses (e.g. 200 OK), preserve local occupancy count and update activity timestamp + synchronized (state) { + state.setLastActivityTime(System.currentTimeMillis()); } - - // Update timestamp for idle TTL lease expiration tracking (35s TTL) + } + } + } else if (msg instanceof LastHttpContent) { + // STEP 2: Release Occupancy on Stream Completion + // When the entire HTTP response payload finishes streaming, decrement the in-flight concurrency count. + PodState state = podRegistry.get(targetWorkerAddress); + if (state != null) { + synchronized (state) { + state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); state.setLastActivityTime(System.currentTimeMillis()); - - if (activeTasksStr != null || leasedNamespace != null) { - LOG.info("shruzard - ProxyBackendHandler: Self-Healed " - + "PodState for {}. Occupancy: {}, Namespace: {}", - targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); - } } } } - - // STEP 2: Relay Worker Response to Client (AppFabric) + + // STEP 3: Relay Worker Response to Client (AppFabric) // Forward the HTTP response header or body chunk directly to the inbound client socket. // Once write completes successfully, request the next chunk from the worker channel. inboundChannel.writeAndFlush(msg).addListener((ChannelFutureListener) future -> { From 8a68d01acce104f15efafbf3adf2ce06314f4d32 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:12:03 +0530 Subject: [PATCH 45/54] refactor(proxy): initialize PodState lastActivityTime cleanly with System.currentTimeMillis() --- .../java/io/cdap/cdap/common/internal/remote/PodState.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java index 70fd4d1d02f5..9f5023deb84c 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -37,9 +37,7 @@ public class PodState { public PodState(String leasedNamespace, int inflightRequests) { this.leasedNamespace = leasedNamespace; this.inflightRequests = inflightRequests; - // Initialize lastActivityTime with a 40s offset so a newly discovered pod is immediately eligible - // to be claimed by any namespace upon startup. - this.lastActivityTime = System.currentTimeMillis() - 40000L; + this.lastActivityTime = System.currentTimeMillis(); } public String getLeasedNamespace() { From bf5a38ed1cd9c9777b981737e065e1dafe2a3033 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:13:22 +0530 Subject: [PATCH 46/54] chore: remove Dockerfile from PR branch --- Dockerfile | 28 ---------------------------- 1 file changed, 28 deletions(-) delete mode 100644 Dockerfile diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index 6ad9384750a2..000000000000 --- a/Dockerfile +++ /dev/null @@ -1,28 +0,0 @@ -# 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 From fd89a7c8d3a4fbcffee9ee8353828517f0ff262a Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:23:00 +0530 Subject: [PATCH 47/54] chore: remove _agents/rules/rbac_taskmanager.md from PR branch --- _agents/rules/rbac_taskmanager.md | 18 ------------------ 1 file changed, 18 deletions(-) delete mode 100644 _agents/rules/rbac_taskmanager.md diff --git a/_agents/rules/rbac_taskmanager.md b/_agents/rules/rbac_taskmanager.md deleted file mode 100644 index ea50b48a302a..000000000000 --- a/_agents/rules/rbac_taskmanager.md +++ /dev/null @@ -1,18 +0,0 @@ -# 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` From 0a2e2237de795c871eca8d164c9e009dec077c10 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:23:17 +0530 Subject: [PATCH 48/54] chore: remove _agents/rules/task_manager_context.md from PR branch --- _agents/rules/task_manager_context.md | 19 ------------------- 1 file changed, 19 deletions(-) delete mode 100644 _agents/rules/task_manager_context.md diff --git a/_agents/rules/task_manager_context.md b/_agents/rules/task_manager_context.md deleted file mode 100644 index 3055a1ae4a36..000000000000 --- a/_agents/rules/task_manager_context.md +++ /dev/null @@ -1,19 +0,0 @@ -# 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. From 70f8bc690d432249f3863f0cf227e1f755ba5c8d Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:28:44 +0530 Subject: [PATCH 49/54] perf(k8s): eliminate blocking synchronous readNamespacedEndpoints call from WatcherThread and rely on asynchronous EndpointsWatcherThread --- .../k8s/discovery/KubeDiscoveryService.java | 35 ++----------------- 1 file changed, 2 insertions(+), 33 deletions(-) diff --git a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java index 0c15b127fcdc..aafe41861bb1 100644 --- a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java +++ b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java @@ -79,6 +79,7 @@ public class KubeDiscoveryService implements DiscoveryService, private static final String SERVICE_TYPE_LOAD_BALANCER = "LoadBalancer"; private static final String SERVICE_TYPE_CLUSTER_IP = "ClusterIP"; private static final String PAYLOAD_NAME = "cdap.service.payload"; + private static final String TASK_WORKER_SERVICE_NAME = "task.worker"; private final String podName; private final String namespace; @@ -223,7 +224,7 @@ public ServiceDiscovered discover(String name) { watcherThread.addService(name); } - if (endpointsWatcherEnabled) { + if (endpointsWatcherEnabled || TASK_WORKER_SERVICE_NAME.equals(name)) { EndpointsWatcherThread endpointsWatcherThread = this.endpointsWatcherThread; if (endpointsWatcherThread == null) { synchronized (this) { @@ -613,38 +614,6 @@ Set toDiscoverables(String name, V1Service service, .orElse(Collections.emptySet()); } - // Try to discover individual Pod IPs via the K8s Endpoints API for ClusterIP services - Set discoverables = new HashSet<>(); - try { - CoreV1Api api = getCoreApi(); - V1Endpoints endpoints = api.readNamespacedEndpoints(meta.getName(), namespace, null); - if (endpoints != null && endpoints.getSubsets() != null) { - for (io.kubernetes.client.openapi.models.V1EndpointSubset subset : endpoints.getSubsets()) { - List addresses = subset.getAddresses(); - List ports = subset.getPorts(); - if (addresses != null && ports != null) { - for (io.kubernetes.client.openapi.models.V1EndpointAddress address : addresses) { - for (io.kubernetes.client.openapi.models.CoreV1EndpointPort port : ports) { - if (servicePorts.stream().anyMatch(sp -> sp.getPort().equals(port.getPort()))) { - Discoverable d = createDiscoverable(name, address.getIp(), - new V1ServicePort().port(port.getPort()), payload); - if (d != null) { - discoverables.add(d); - } - } - } - } - } - } - } - } catch (Exception e) { - LOG.warn("Failed to retrieve endpoints for service {}, falling back to service hostname", name, e); - } - - if (!discoverables.isEmpty()) { - return discoverables; - } - String hostname = String.format("%s.%s", meta.getName(), namespace); return servicePorts.stream() .map(port -> createDiscoverable( From 27c7b4594e0b448b7eaa604e01d3de5ddc134994 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:50:09 +0530 Subject: [PATCH 50/54] refactor(k8s): simplify EndpointsWatcherThread start check to only check endpointsWatcherEnabled --- .../java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java index aafe41861bb1..dbab86419f69 100644 --- a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java +++ b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/discovery/KubeDiscoveryService.java @@ -79,7 +79,6 @@ public class KubeDiscoveryService implements DiscoveryService, private static final String SERVICE_TYPE_LOAD_BALANCER = "LoadBalancer"; private static final String SERVICE_TYPE_CLUSTER_IP = "ClusterIP"; private static final String PAYLOAD_NAME = "cdap.service.payload"; - private static final String TASK_WORKER_SERVICE_NAME = "task.worker"; private final String podName; private final String namespace; @@ -224,7 +223,7 @@ public ServiceDiscovered discover(String name) { watcherThread.addService(name); } - if (endpointsWatcherEnabled || TASK_WORKER_SERVICE_NAME.equals(name)) { + if (endpointsWatcherEnabled) { EndpointsWatcherThread endpointsWatcherThread = this.endpointsWatcherThread; if (endpointsWatcherThread == null) { synchronized (this) { From 4d9c30101236b7bb8241beaefc150fcc66b1fa45 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Wed, 19 Aug 2026 11:48:06 +0000 Subject: [PATCH 51/54] refactor: rename TaskManagerMain to TaskManagerServiceMain --- .../k8s/{TaskManagerMain.java => TaskManagerServiceMain.java} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/{TaskManagerMain.java => TaskManagerServiceMain.java} (95%) diff --git a/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerServiceMain.java similarity index 95% rename from cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java rename to cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerServiceMain.java index 8cc815eddee5..b8d3bd7b177c 100644 --- a/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerMain.java +++ b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerServiceMain.java @@ -46,10 +46,10 @@ *
  • Wires standard CDAP logging context under {@code task-manager} for Cloud Logging.
    • *
*/ -public class TaskManagerMain extends AbstractServiceMain { +public class TaskManagerServiceMain extends AbstractServiceMain { public static void main(String[] args) throws Exception { - main(TaskManagerMain.class, args); + main(TaskManagerServiceMain.class, args); } @Override From dd5e7fbd1e28940b868a8cb97a7f43e596bb23e3 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Thu, 20 Aug 2026 05:59:18 +0000 Subject: [PATCH 52/54] make TaskManager register itself to KubeDiscoverService --- .../internal/remote/TaskManagerService.java | 20 ++++++++++++++++++- task-manager-service.yaml | 2 +- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index 63e39f369365..fe0a6038f701 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -33,6 +33,11 @@ import org.apache.twill.discovery.Discoverable; import org.apache.twill.discovery.DiscoveryServiceClient; +import org.apache.twill.discovery.DiscoveryService; +import org.apache.twill.common.Cancellable; +import io.cdap.cdap.common.discovery.URIScheme; +import io.cdap.cdap.common.discovery.ResolvingDiscoverable; +import java.net.InetSocketAddress; import io.cdap.cdap.common.conf.Constants; import java.util.Set; import java.util.HashSet; @@ -69,12 +74,15 @@ public class TaskManagerService extends AbstractIdleService { private final Map podRegistry = new ConcurrentHashMap<>(); private final DiscoveryServiceClient discoveryServiceClient; + private final DiscoveryService discoveryService; + private Cancellable cancellable; @Inject - TaskManagerService(CConfiguration cConf, DiscoveryServiceClient discoveryServiceClient) { + TaskManagerService(CConfiguration cConf, DiscoveryServiceClient discoveryServiceClient, DiscoveryService discoveryService) { this.port = cConf.getInt("task.manager.port", 11025); this.address = cConf.get("task.manager.address", "0.0.0.0"); this.discoveryServiceClient = discoveryServiceClient; + this.discoveryService = discoveryService; LOG.info("shruzard - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); } @@ -104,12 +112,22 @@ protected void initChannel(SocketChannel ch) { }); channelFuture = b.bind(address, port).sync(); + + // Announce via DiscoveryService so Kubernetes Discovery can provision K8s Service/Endpoints dynamically + InetSocketAddress socketAddress = new InetSocketAddress(address, port); + this.cancellable = discoveryService.register( + ResolvingDiscoverable.of(URIScheme.HTTP.createDiscoverable(Constants.Service.TASK_MANAGER, socketAddress)) + ); + LOG.info("shruzard - TaskManagerService Proxy HTTP server started successfully at {}:{}", address, port); } @Override protected void shutDown() throws Exception { LOG.info("shruzard - Stopping TaskManagerService Proxy HTTP server..."); + if (this.cancellable != null) { + this.cancellable.cancel(); + } if (channelFuture != null) { channelFuture.channel().close().sync(); } diff --git a/task-manager-service.yaml b/task-manager-service.yaml index b9a939d7c31a..a1036c83a1ac 100644 --- a/task-manager-service.yaml +++ b/task-manager-service.yaml @@ -37,7 +37,7 @@ spec: image: us-east1-docker.pkg.dev/j145774183a931adb-tp/cdf-dev-shru/cloud-data-fusion:latest imagePullPolicy: Always args: - - "io.cdap.cdap.master.environment.k8s.TaskManagerMain" + - "io.cdap.cdap.master.environment.k8s.TaskManagerServiceMain" - "--env=k8s" env: - name: SERVICE_NAME From 2eb474123aeac0edc0b059d68760817cb6f5ef86 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 25 Aug 2026 09:24:47 +0000 Subject: [PATCH 53/54] update PodState to use AtomicReference and CAS mechanism for lock free synchronization --- .../cdap/common/internal/remote/PodState.java | 120 ++++++++++++++---- .../internal/remote/ProxyBackendHandler.java | 61 ++++----- .../internal/remote/ProxyFrontendHandler.java | 67 +++++----- .../common/internal/remote/RemoteClient.java | 48 +++---- .../internal/remote/RemoteTaskExecutor.java | 42 +++++- .../internal/remote/TaskManagerService.java | 53 +++++--- 6 files changed, 253 insertions(+), 138 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java index 9f5023deb84c..77df37e40e2d 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -16,51 +16,117 @@ package io.cdap.cdap.common.internal.remote; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + /** * PodState represents the in-memory routing and lease status of an individual Task Worker pod. - * - *

It tracks: - *

    - *
  • {@code leasedNamespace}: The namespace currently pinned to this physical worker pod. - * Only requests belonging to this namespace may execute on this pod.
    • - *
  • {@code inflightRequests}: The number of active concurrent tasks running on this pod - * (governed up to 10 concurrent requests).
    • - *
  • {@code lastActivityTime}: Timestamp of the most recent request completion, used to calculate - * idle TTL eviction (35s) so idle pods can be reclaimed by other namespaces.
    • - *
+ * Entirely lock-free, backing state via an immutable internal representation and AtomicReference CAS loops. */ public class PodState { - private String leasedNamespace; - private int inflightRequests; - private long lastActivityTime; + private static class State { + final String leasedNamespace; + final int inflightRequests; + final long lastActivityTime; + + State(String leasedNamespace, int inflightRequests, long lastActivityTime) { + this.leasedNamespace = leasedNamespace; + this.inflightRequests = inflightRequests; + this.lastActivityTime = lastActivityTime; + } + } + + private final AtomicReference stateRef; public PodState(String leasedNamespace, int inflightRequests) { - this.leasedNamespace = leasedNamespace; - this.inflightRequests = inflightRequests; - this.lastActivityTime = System.currentTimeMillis(); + this.stateRef = new AtomicReference<>(new State( + leasedNamespace, + inflightRequests, + System.nanoTime() - TimeUnit.SECONDS.toNanos(40) + )); } public String getLeasedNamespace() { - return leasedNamespace; + return stateRef.get().leasedNamespace; + } + + public int getInflightRequests() { + return stateRef.get().inflightRequests; } - public void setLeasedNamespace(String leasedNamespace) { - this.leasedNamespace = leasedNamespace; + public long getLastActivityTime() { + return stateRef.get().lastActivityTime; } - public int getInflightRequests() { - return inflightRequests; + public boolean tryAcquireWarmLease(String namespace, int maxConcurrency) { + while (true) { + State current = stateRef.get(); + if (!namespace.equals(current.leasedNamespace) || current.inflightRequests >= maxConcurrency) { + return false; + } + State next = new State(current.leasedNamespace, current.inflightRequests + 1, System.nanoTime()); + if (stateRef.compareAndSet(current, next)) { + return true; + } + } } - public void setInflightRequests(int inflightRequests) { - this.inflightRequests = inflightRequests; + public boolean tryClaimIdleLease(String namespace, long idleTimeoutNanos) { + while (true) { + State current = stateRef.get(); + boolean isUnleased = current.leasedNamespace == null || current.leasedNamespace.isEmpty(); + boolean isExpiredIdle = current.inflightRequests == 0 + && (System.nanoTime() - current.lastActivityTime > idleTimeoutNanos); + + if (current.inflightRequests != 0 || (!isUnleased && !isExpiredIdle)) { + return false; + } + State next = new State(namespace, 1, System.nanoTime()); + if (stateRef.compareAndSet(current, next)) { + return true; + } + } } - public long getLastActivityTime() { - return lastActivityTime; + public void decrementInflightRequests() { + while (true) { + State current = stateRef.get(); + State next = new State(current.leasedNamespace, + Math.max(0, current.inflightRequests - 1), System.nanoTime()); + if (stateRef.compareAndSet(current, next)) { + return; + } + } } - public void setLastActivityTime(long lastActivityTime) { - this.lastActivityTime = lastActivityTime; + public void updateFromHeader(String activeTasksStr, String leasedNamespace) { + while (true) { + State current = stateRef.get(); + int nextInflight = current.inflightRequests; + if (activeTasksStr != null) { + try { + nextInflight = Integer.parseInt(activeTasksStr); + } catch (NumberFormatException e) { + nextInflight = Math.max(0, current.inflightRequests - 1); + } + } else { + nextInflight = Math.max(0, current.inflightRequests - 1); + } + String nextNamespace = leasedNamespace != null ? leasedNamespace : current.leasedNamespace; + State next = new State(nextNamespace, nextInflight, System.nanoTime()); + if (stateRef.compareAndSet(current, next)) { + return; + } + } + } + + public void recordActivity() { + while (true) { + State current = stateRef.get(); + State next = new State(current.leasedNamespace, current.inflightRequests, System.nanoTime()); + if (stateRef.compareAndSet(current, next)) { + return; + } + } } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index 7a2b17560ab8..bf9f2c77f333 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -57,6 +57,8 @@ public class ProxyBackendHandler extends ChannelInboundHandlerAdapter { private final Map podRegistry; private final String targetWorkerAddress; + private boolean decremented = false; + public ProxyBackendHandler(Channel inboundChannel, Map podRegistry, String targetWorkerAddress) { this.inboundChannel = inboundChannel; this.podRegistry = podRegistry; @@ -73,50 +75,27 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { if (state != null) { // STEP 1: Selective Self-Healing & Occupancy Synchronization // ONLY synchronize ground truth from headers when the worker explicitly rejects the request - // with 409 Conflict (namespace mismatch / split brain) or 429 Too Many Requests (worker saturated). if (statusCode == HttpResponseStatus.CONFLICT.code() || statusCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - synchronized (state) { - String activeTasksStr = resp.headers().get("X-Active-Tasks"); - String leasedNamespace = resp.headers().get("X-Leased-Namespace"); - - if (activeTasksStr != null) { - try { - state.setInflightRequests(Integer.parseInt(activeTasksStr)); - } catch (NumberFormatException e) { - state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); - } - } else { - state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); - } - - if (leasedNamespace != null) { - state.setLeasedNamespace(leasedNamespace); - } - - state.setLastActivityTime(System.currentTimeMillis()); - - LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState after status {} for {}. " - + "Occupancy: {}, Namespace: {}", - statusCode, targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); - } + + String activeTasksStr = resp.headers().get("X-Active-Tasks"); + String leasedNamespace = resp.headers().get("X-Leased-Namespace"); + + state.updateFromHeader(activeTasksStr, leasedNamespace); + decremented = true; + + LOG.info("shruzard - ProxyBackendHandler: Self-Healed PodState after status {} for {}. " + + "Occupancy: {}, Namespace: {}", + statusCode, targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); } else { // For normal responses (e.g. 200 OK), preserve local occupancy count and update activity timestamp - synchronized (state) { - state.setLastActivityTime(System.currentTimeMillis()); - } + state.recordActivity(); } } } else if (msg instanceof LastHttpContent) { // STEP 2: Release Occupancy on Stream Completion // When the entire HTTP response payload finishes streaming, decrement the in-flight concurrency count. - PodState state = podRegistry.get(targetWorkerAddress); - if (state != null) { - synchronized (state) { - state.setInflightRequests(Math.max(0, state.getInflightRequests() - 1)); - state.setLastActivityTime(System.currentTimeMillis()); - } - } + releaseOccupancy(); } // STEP 3: Relay Worker Response to Client (AppFabric) @@ -142,8 +121,19 @@ public void channelWritabilityChanged(ChannelHandlerContext ctx) { ctx.fireChannelWritabilityChanged(); } + private void releaseOccupancy() { + if (!decremented) { + PodState state = podRegistry.get(targetWorkerAddress); + if (state != null) { + state.decrementInflightRequests(); + } + decremented = true; + } + } + @Override public void channelInactive(ChannelHandlerContext ctx) { + releaseOccupancy(); // If backend worker disconnects or crashes, flush and close the client socket ProxyFrontendHandler.closeOnFlush(inboundChannel); } @@ -151,6 +141,7 @@ public void channelInactive(ChannelHandlerContext ctx) { @Override public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { cause.printStackTrace(); + releaseOccupancy(); ProxyFrontendHandler.closeOnFlush(ctx.channel()); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index aa1b3a7dbd25..33e0c7d4a45e 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -87,6 +87,7 @@ public ProxyFrontendHandler(Map podRegistry, DiscoveryServiceC @Override public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { if (msg instanceof HttpRequest) { + LOG.info("shruzard - ProxyFrontendHandler Received request!"); HttpRequest req = (HttpRequest) msg; // STEP 0: Discover Live Task Worker Pods (Zero-Stale State) @@ -105,49 +106,52 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception } podRegistry.keySet().removeIf(existingPod -> !activePods.contains(existingPod)); + LOG.info("shruzard - ProxyFrontendHandler leases: [{}]", + podRegistry.entrySet().stream() + .map(e -> e.getKey() + "=" + + (e.getValue().getLeasedNamespace() == null + ? "null" : e.getValue().getLeasedNamespace() + "_" + e.getValue().getInflightRequests())) + .collect(java.util.stream.Collectors.joining(", "))); + // Extract target namespace from the request header (defaults to "default" if omitted) String targetNamespace = req.headers().get("X-CDF-Namespace"); if (targetNamespace == null) targetNamespace = "default"; String targetWorkerAddress = null; + // We deliberately evaluate sequentially (no randomization/shuffling) to maximize pod density. + // This ensures a namespace fills a single pod to its maximum capacity (10 tasks) before + // spilling over and leasing an entirely new empty pod, preserving cluster availability + // for other namespaces. + // STEP 1: Warm Match Selection - // Look for a pod already leased to this exact namespace that has capacity (< 10 inflight tasks). - // Reusing warm pods avoids expensive Workload Identity / ArtifactLocalizer re-authentication. + // Perform lock-free Compare-And-Swap evaluation using the AtomicReference loop for (Map.Entry entry : podRegistry.entrySet()) { PodState state = entry.getValue(); - synchronized (state) { - if (targetNamespace.equals(state.getLeasedNamespace()) && state.getInflightRequests() < 10) { - targetWorkerAddress = entry.getKey(); - state.setInflightRequests(state.getInflightRequests() + 1); - LOG.info("shruzard - ProxyFrontendHandler: Found warm match " - + "for '{}' at {}. Occupancy: {}", - targetNamespace, targetWorkerAddress, state.getInflightRequests()); - break; - } + + if (state.tryAcquireWarmLease(targetNamespace, 10)) { + targetWorkerAddress = entry.getKey(); + LOG.info("shruzard - ProxyFrontendHandler: Found warm match " + + "for '{}' at {}. Occupancy: {}", + targetNamespace, targetWorkerAddress, state.getInflightRequests()); + break; } } // STEP 2: Idle Pod Claiming - // If no warm pod has capacity, claim an unleased pod or an idle pod whose lease expired (35s TTL). if (targetWorkerAddress == null) { + // Determine our starvation/eviction threshold (35 seconds in NanoTime) + long idleTimeoutNanos = java.util.concurrent.TimeUnit.SECONDS.toNanos(35); + for (Map.Entry entry : podRegistry.entrySet()) { PodState state = entry.getValue(); - synchronized (state) { - boolean isUnleased = (state.getLeasedNamespace() == null - || state.getLeasedNamespace().isEmpty()); - boolean isExpiredIdle = (state.getInflightRequests() == 0 - && (System.currentTimeMillis() - state.getLastActivityTime() > 35000)); - - if (state.getInflightRequests() == 0 && (isUnleased || isExpiredIdle)) { - targetWorkerAddress = entry.getKey(); - state.setLeasedNamespace(targetNamespace); - state.setInflightRequests(1); - LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " - + "(Unleased: {}, ExpiredIdle: {}) at {} for namespace '{}'.", - isUnleased, isExpiredIdle, targetWorkerAddress, targetNamespace); - break; - } + + if (state.tryClaimIdleLease(targetNamespace, idleTimeoutNanos)) { + targetWorkerAddress = entry.getKey(); + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + + "for new namespace '{}' at {}. Previous occupant evicted.", + targetNamespace, targetWorkerAddress); + break; } } } @@ -170,6 +174,9 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception final String chosenWorker = targetWorkerAddress; String[] hostPort = targetWorkerAddress.split(":"); + LOG.info("shruzard - ProxyFrontendHandler: Setting up TCP connection to task worker IP - {} namespace {}", + targetWorkerAddress, targetNamespace); + // STEP 4: Establish Outbound TCP Socket to Chosen Task Worker Pod // 1. Temporarily pause reading from the client (AppFabric) socket so data does not pile up in RAM // while the TCP handshake to the worker is completing. @@ -179,6 +186,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception // 2. Initialize the outbound Netty client Bootstrap. // Sharing ctx.channel().eventLoop() ensures that both inbound and outbound channels run on the same // event loop thread, guaranteeing thread safety without thread context-switching overhead. + Bootstrap b = new Bootstrap(); b.group(ctx.channel().eventLoop()) .channel(NioSocketChannel.class) @@ -211,6 +219,7 @@ protected void initChannel(SocketChannel ch) { connecting = false; if (future.isSuccess()) { // Flush any request headers/chunks that arrived while TCP connection was being negotiated + LOG.info("shruzard - ProxyFrontendHandler: Connected with task worker successfully! "); Object pendingMsg = pendingMessages.poll(); while (pendingMsg != null) { outboundChannel.write(pendingMsg); @@ -232,9 +241,7 @@ protected void initChannel(SocketChannel ch) { // Decrement inflight count on failed connection PodState fallbackState = podRegistry.get(chosenWorker); if (fallbackState != null) { - synchronized (fallbackState) { - fallbackState.setInflightRequests(Math.max(0, fallbackState.getInflightRequests() - 1)); - } + fallbackState.decrementInflightRequests(); } ctx.channel().close(); } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index d89dda4c2b50..17d9b0e30c42 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java @@ -69,7 +69,6 @@ public class RemoteClient { private static final Logger LOG = LoggerFactory.getLogger(RemoteClient.class); - private static final String TASK_MANAGER_URL = "http://cdap-task-manager.default.svc.cluster.local:11025"; private static final Gson GSON = new Gson(); private final InternalAuthenticator internalAuthenticator; @@ -95,12 +94,20 @@ public class RemoteClient { this.discoverableServiceName = discoverableServiceName; this.httpRequestConfig = httpRequestConfig; this.discoveryClient = discoveryClient; + if (Constants.Service.TASK_MANAGER.equals(discoverableServiceName) + || Constants.Service.TASK_WORKER.equals(discoverableServiceName)) { + LOG.info("shruzard - RemoteClient L97 discoverableServiceName- {}", + discoverableServiceName); + } + this.endpointStrategy = new RandomEndpointStrategy( () -> discoveryClient.discover(discoverableServiceName)); String cleanBasePath = basePath.startsWith("/") ? basePath.substring(1) : basePath; this.basePath = cleanBasePath.endsWith("/") ? cleanBasePath : cleanBasePath + "/"; this.remoteAuthenticator = remoteAuthenticator; + this.rbacEnabled = rbacEnabled; + } /** @@ -301,31 +308,28 @@ public URL resolve(String resource) { * null, it falls back to the default random discovery strategy. */ public URL resolve(String resource, @Nullable String routingKey) { - if (!rbacEnabled || routingKey == null || (!Constants.Service.TASK_MANAGER.equals(discoverableServiceName) - && !Constants.Service.TASK_WORKER.equals(discoverableServiceName))) { - Discoverable discoverable = endpointStrategy.pick(1L, TimeUnit.SECONDS); - if (discoverable == null) { - throw new ServiceUnavailableException(discoverableServiceName); - } - if(!rbacEnabled) { - LOG.info("shruzard - RemoteClient RBAC disabled not using Task Manager", routingKey); - } - URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); - try { - return rewriteUrl(uri.toURL()); - } catch (MalformedURLException e) { - throw new IllegalStateException( - String.format("Discovered service %s, but it announced malformed URL %s", - discoverableServiceName, uri), e); - } + Discoverable discoverable = endpointStrategy.pick(1L, TimeUnit.SECONDS); + if (discoverable == null) { + throw new ServiceUnavailableException(discoverableServiceName); + } + if(discoverableServiceName.equals(Constants.Service.TASK_MANAGER) && rbacEnabled) { + LOG.info("shruzard AA RemoteClient RBAC enabled dynamically discovered {} via Twill/K8s for routingKey: {}", + discoverableServiceName, routingKey); } - LOG.info("shruzard - RemoteClient routing directly to Netty TaskManager L7 proxy for routingKey: {}", routingKey); + URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); try { - String cleanPath = (basePath + resource).replaceAll("//+", "/"); - return new URL(TASK_MANAGER_URL + "/" + cleanPath); + if (Constants.Service.TASK_MANAGER.equals(discoverableServiceName) + || Constants.Service.TASK_WORKER.equals(discoverableServiceName)) { + LOG.info("shruzard - RemoteClient L313 calling task worker with url: {}", + rewriteUrl(uri.toURL())); + } + + return rewriteUrl(uri.toURL()); } catch (MalformedURLException e) { - throw new ServiceUnavailableException(discoverableServiceName, e); + throw new IllegalStateException( + String.format("Discovered service %s, but it announced malformed URL %s", + discoverableServiceName, uri), e); } } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index cd6bb5f4ecd0..1f7647938754 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -83,6 +83,7 @@ public class RemoteTaskExecutor { || (throwable instanceof NoRouteToHostException); private final boolean compression; private final RemoteClient remoteClient; + private final RemoteClient fallbackClient; private final RetryStrategy retryStrategy; private final Predicate retryablePredicate; private final MetricsCollectionService metricsCollectionService; @@ -105,9 +106,19 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics ? Constants.Service.TASK_MANAGER : Constants.Service.TASK_WORKER; String serviceName = workerType == Type.TASK_WORKER ? taskServiceName : Constants.Service.SYSTEM_WORKER; + LOG.info("shruzard - RemoteTaskExecutor: Using serviceName - {}", + serviceName); + this.remoteClient = remoteClientFactory.createRemoteClient(serviceName, httpRequestConfig, Constants.Gateway.INTERNAL_API_VERSION_3); + + // Explicitly scope the fallback client directly to TASK_WORKER bypass K8s Service Proxy IPs completely + // when the 60s Circuit Breaker activates `routingKey = null`. + this.fallbackClient = remoteClientFactory.createRemoteClient( + workerType == Type.TASK_WORKER ? Constants.Service.TASK_WORKER : Constants.Service.SYSTEM_WORKER, + httpRequestConfig, Constants.Gateway.INTERNAL_API_VERSION_3); + this.metricsCollectionService = metricsCollectionService; this.userEncryptionAeadCipher = aeadCipher; this.fallbackTimeoutMs = 60000; @@ -170,7 +181,12 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception // STEP 3: Construct Outbound HTTP Request with Namespace Header // Inject X-CDF-Namespace so the Netty Proxy can route this request to a warm pod leased to this namespace. - HttpRequest.Builder requestBuilder = remoteClient + LOG.info("shruzard - RemoteTaskExecutor: Sending request to RemoteClient with routingKey: {}", + routingKey); + + RemoteClient activeClient = routingKey == null ? fallbackClient : remoteClient; + + HttpRequest.Builder requestBuilder = activeClient .requestBuilder(HttpMethod.POST, workerUrl, routingKey) .addHeader("X-CDF-Namespace", namespace) .withBody(requestBody.duplicate()); @@ -189,7 +205,7 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } HttpRequest httpRequest = requestBuilder.build(); - HttpResponse httpResponse = remoteClient.execute(httpRequest); + HttpResponse httpResponse = activeClient.execute(httpRequest); proxyReachable.set(true); // Resetting user credentials for further execution of current request @@ -202,9 +218,9 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception // so CDAP's Retries framework retries with exponential backoff until a lease slot opens. if (httpResponse.getResponseCode() == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { throw new RetryableException( - String.format("Task Worker cluster is fully saturated. " - + "Unable to secure a compute lease after " - + "60 seconds (HTTP 429 for %s). Please try again.", + String.format("Task Worker cluster is fully saturated (HTTP 429). " + + "The Proxy could not immediately secure a compute lease for %s. " + + "Triggering exponential backoff...", runnableTaskRequest.getClassName())); } if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { @@ -220,6 +236,17 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception throw new RetryableException( String.format("Received exception %s for %s", e.getMessage(), runnableTaskRequest.getClassName())); + } catch (ServiceException e) { + // 503 natively throws ServiceUnavailableException (which extends RetryableException). + // But 502 and 504 throw plain ServiceException which would bypass our circuit breaker and crash! + // We manually trap 502/504 infrastructure errors during POST and force them to retry, + // ensuring they loop for 60 seconds and correctly trigger the TaskWorker fallback bypass. + if (e.getStatusCode() == HttpResponseStatus.BAD_GATEWAY.code() || + e.getStatusCode() == HttpResponseStatus.GATEWAY_TIMEOUT.code()) { + throw new RetryableException("Proxy infrastructure unreachable (HTTP " + e.getStatusCode() + + "). Forcing retry to trigger Circuit Breaker.", e); + } + throw e; // Non-infrastructure ServiceExceptions (like 403 or 401) must fail immediately } }, retryStrategy, retryablePredicate); } catch (ServiceException se) { @@ -233,7 +260,10 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception if (e instanceof RetryableException && e.getMessage() != null && e.getMessage().contains("Task Worker cluster is fully saturated")) { throw new ServiceException( - e.getMessage(), e, HttpResponseStatus.TOO_MANY_REQUESTS); + String.format("Task Worker cluster is fully saturated. " + + "Unable to secure a compute lease in 60 seconds (HTTP 429). " + + "Please try again later."), + e, HttpResponseStatus.TOO_MANY_REQUESTS); } throw e; } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java index fe0a6038f701..fe87a172340c 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -16,9 +16,23 @@ package io.cdap.cdap.common.internal.remote; +import java.net.InetSocketAddress; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import org.apache.twill.common.Cancellable; +import org.apache.twill.discovery.DiscoveryService; +import org.apache.twill.discovery.DiscoveryServiceClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import com.google.common.util.concurrent.AbstractIdleService; import com.google.inject.Inject; + import io.cdap.cdap.common.conf.CConfiguration; +import io.cdap.cdap.common.conf.Constants; +import io.cdap.cdap.common.discovery.ResolvingDiscoverable; +import io.cdap.cdap.common.discovery.URIScheme; import io.netty.bootstrap.ServerBootstrap; import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelInitializer; @@ -28,21 +42,6 @@ import io.netty.channel.socket.SocketChannel; import io.netty.channel.socket.nio.NioServerSocketChannel; import io.netty.handler.codec.http.HttpServerCodec; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import org.apache.twill.discovery.Discoverable; -import org.apache.twill.discovery.DiscoveryServiceClient; -import org.apache.twill.discovery.DiscoveryService; -import org.apache.twill.common.Cancellable; -import io.cdap.cdap.common.discovery.URIScheme; -import io.cdap.cdap.common.discovery.ResolvingDiscoverable; -import java.net.InetSocketAddress; -import io.cdap.cdap.common.conf.Constants; -import java.util.Set; -import java.util.HashSet; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; /** * TaskManagerService runs the Centralized Netty Proxy server inside the Task Manager pod. @@ -78,9 +77,10 @@ public class TaskManagerService extends AbstractIdleService { private Cancellable cancellable; @Inject - TaskManagerService(CConfiguration cConf, DiscoveryServiceClient discoveryServiceClient, DiscoveryService discoveryService) { - this.port = cConf.getInt("task.manager.port", 11025); - this.address = cConf.get("task.manager.address", "0.0.0.0"); + TaskManagerService(CConfiguration cConf, DiscoveryServiceClient discoveryServiceClient, + DiscoveryService discoveryService) { + this.port = cConf.getInt("task.manager.bind.port", 11025); + this.address = cConf.get("task.manager.bind.address", "0.0.0.0"); this.discoveryServiceClient = discoveryServiceClient; this.discoveryService = discoveryService; @@ -91,6 +91,23 @@ public class TaskManagerService extends AbstractIdleService { protected void startUp() throws Exception { LOG.info("shruzard - Starting TaskManagerService Proxy HTTP server..."); + // Enable live Kubernetes Endpoints streaming if supported by the discovery client, + // guaranteeing we get direct Pod IPs (V1Endpoints) instead of load-balanced ClusterIPs (V1Service). + try { + java.lang.reflect.Method method = discoveryServiceClient.getClass().getMethod("enableEndpointsWatcher"); + method.invoke(discoveryServiceClient); + LOG.info("shruzard - Successfully enabled Kubernetes Endpoints watcher on discovery service {}", + discoveryServiceClient.getClass().getSimpleName()); + } catch (NoSuchMethodException ignored) { + // Normal for discovery services that do not support endpoints watching (e.g. In-memory, ZK) + } catch (Exception e) { + LOG.warn("shruzard - Failed to invoke enableEndpointsWatcher on discovery service", e); + } + + // Pre-warm the Twill discovery cache asynchronously to ensure K8s Watchers + // are fully hydrated before the first HTTP request hits the proxy. + discoveryServiceClient.discover(Constants.Service.TASK_WORKER); + bossGroup = new NioEventLoopGroup(1, new com.google.common.util.concurrent.ThreadFactoryBuilder() .setNameFormat("taskmanager-boss-thread-%d").build()); From 3482e2383074c98ec4f105e8a010b6553fa262f2 Mon Sep 17 00:00:00 2001 From: shruthi713 Date: Tue, 25 Aug 2026 15:25:58 +0000 Subject: [PATCH 54/54] update logs --- .../internal/remote/ProxyBackendHandler.java | 12 ++++++++++- .../internal/remote/ProxyFrontendHandler.java | 2 +- .../internal/remote/RemoteTaskExecutor.java | 21 +++++++++++++++---- 3 files changed, 29 insertions(+), 6 deletions(-) diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java index bf9f2c77f333..a9c4e53ec987 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -70,6 +70,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { if (msg instanceof HttpResponse) { HttpResponse resp = (HttpResponse) msg; int statusCode = resp.status().code(); + LOG.info("shruzard - ProxyBackendHandler: Received status code : {} from task worker", statusCode); + PodState state = podRegistry.get(targetWorkerAddress); if (state != null) { @@ -77,7 +79,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { // ONLY synchronize ground truth from headers when the worker explicitly rejects the request if (statusCode == HttpResponseStatus.CONFLICT.code() || statusCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - + String activeTasksStr = resp.headers().get("X-Active-Tasks"); String leasedNamespace = resp.headers().get("X-Leased-Namespace"); @@ -89,12 +91,18 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { statusCode, targetWorkerAddress, state.getInflightRequests(), state.getLeasedNamespace()); } else { // For normal responses (e.g. 200 OK), preserve local occupancy count and update activity timestamp + LOG.info("shruzard - ProxyBackendHandler: Received! status code : {} " + + "from task worker. Updating activityTimestamp", statusCode); + state.recordActivity(); } } } else if (msg instanceof LastHttpContent) { // STEP 2: Release Occupancy on Stream Completion // When the entire HTTP response payload finishes streaming, decrement the in-flight concurrency count. + LOG.info("shruzard - ProxyBackendHandler: Received last message from TaskWorker. " + + "Releasing occupancy in map.. "); + releaseOccupancy(); } @@ -105,6 +113,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { if (future.isSuccess()) { ctx.channel().read(); } else { + LOG.info("shruzard - ProxyBackendHandler: Unable to write back to App fabric... Closing channel "); + future.channel().close(); } }); diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java index 33e0c7d4a45e..8a3c8fdcd649 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -149,7 +149,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (state.tryClaimIdleLease(targetNamespace, idleTimeoutNanos)) { targetWorkerAddress = entry.getKey(); LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " - + "for new namespace '{}' at {}. Previous occupant evicted.", + + "for new namespace '{}' at {}", targetNamespace, targetWorkerAddress); break; } diff --git a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java index 1f7647938754..1fda55b5117f 100644 --- a/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteTaskExecutor.java @@ -91,6 +91,7 @@ public class RemoteTaskExecutor { private final String workerUrl; private final boolean isWorkerEncryptionRequired; private final long fallbackTimeoutMs; + private final String serviceName; public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metricsCollectionService, RemoteClientFactory remoteClientFactory, Type workerType, AeadCipher aeadCipher) { @@ -104,7 +105,7 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics this.compression = cConf.getBoolean(Constants.TaskWorker.COMPRESSION_ENABLED); String taskServiceName = cConf.getBoolean(Constants.Security.Authorization.ENABLED) ? Constants.Service.TASK_MANAGER : Constants.Service.TASK_WORKER; - String serviceName = workerType == Type.TASK_WORKER + this.serviceName = workerType == Type.TASK_WORKER ? taskServiceName : Constants.Service.SYSTEM_WORKER; LOG.info("shruzard - RemoteTaskExecutor: Using serviceName - {}", serviceName); @@ -112,7 +113,8 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics this.remoteClient = remoteClientFactory.createRemoteClient(serviceName, httpRequestConfig, Constants.Gateway.INTERNAL_API_VERSION_3); - + LOG.info("shruzard - RemoteTaskExecutor: Creating fallbackClient " + + "to bypass proxy during failure"); // Explicitly scope the fallback client directly to TASK_WORKER bypass K8s Service Proxy IPs completely // when the 60s Circuit Breaker activates `routingKey = null`. this.fallbackClient = remoteClientFactory.createRemoteClient( @@ -205,13 +207,23 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } HttpRequest httpRequest = requestBuilder.build(); + + long requestStartTime = System.currentTimeMillis(); HttpResponse httpResponse = activeClient.execute(httpRequest); + long requestEndTime = System.currentTimeMillis(); + long executionDurationMs = requestEndTime - requestStartTime; + proxyReachable.set(true); // Resetting user credentials for further execution of current request if (isWorkerEncryptionRequired) { SecurityRequestContext.setUserCredential(currentCredential); } + LOG.info("shruzard - RemoteTaskExecutor: Received response from " + + "{} with status code {} in {} ms", + this.serviceName, + httpResponse.getResponseCode(), + executionDurationMs); // STEP 4: Handle Responses & Retryable Exceptions // If proxy/worker returned 429 Too Many Requests (cluster saturated), throw RetryableException @@ -231,6 +243,7 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception byte[] result = httpResponse.getUncompressedResponseBody(); //emit metrics with successful result emitMetrics(startTime, true, runnableTaskRequest, retryContext.getRetryAttempt()); + return result; } catch (NoRouteToHostException e) { throw new RetryableException( @@ -241,8 +254,8 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception // But 502 and 504 throw plain ServiceException which would bypass our circuit breaker and crash! // We manually trap 502/504 infrastructure errors during POST and force them to retry, // ensuring they loop for 60 seconds and correctly trigger the TaskWorker fallback bypass. - if (e.getStatusCode() == HttpResponseStatus.BAD_GATEWAY.code() || - e.getStatusCode() == HttpResponseStatus.GATEWAY_TIMEOUT.code()) { + if (e.getStatusCode() == HttpResponseStatus.BAD_GATEWAY.code() + || e.getStatusCode() == HttpResponseStatus.GATEWAY_TIMEOUT.code()) { throw new RetryableException("Proxy infrastructure unreachable (HTTP " + e.getStatusCode() + "). Forcing retry to trigger Circuit Breaker.", e); }