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/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/PodState.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java new file mode 100644 index 000000000000..77df37e40e2d --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/PodState.java @@ -0,0 +1,132 @@ +/* + * 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.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +/** + * PodState represents the in-memory routing and lease status of an individual Task Worker pod. + * Entirely lock-free, backing state via an immutable internal representation and AtomicReference CAS loops. + */ +public class PodState { + 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.stateRef = new AtomicReference<>(new State( + leasedNamespace, + inflightRequests, + System.nanoTime() - TimeUnit.SECONDS.toNanos(40) + )); + } + + public String getLeasedNamespace() { + return stateRef.get().leasedNamespace; + } + + public int getInflightRequests() { + return stateRef.get().inflightRequests; + } + + public long getLastActivityTime() { + return stateRef.get().lastActivityTime; + } + + 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 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 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 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 new file mode 100644 index 000000000000..a9c4e53ec987 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyBackendHandler.java @@ -0,0 +1,157 @@ +/* + * 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 io.netty.handler.codec.http.HttpResponseStatus; +import io.netty.handler.codec.http.LastHttpContent; + +import java.util.Map; +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 upon rejections, and relays the HTTP response bytes + * directly back to the inbound client (AppFabric). + * + *

Key Responsibilities: + *

+ */ +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; + this.podRegistry = podRegistry; + this.targetWorkerAddress = targetWorkerAddress; + } + + @Override + 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) { + // STEP 1: Selective Self-Healing & Occupancy Synchronization + // 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"); + + 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 + 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(); + } + + // 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 -> { + if (future.isSuccess()) { + ctx.channel().read(); + } else { + LOG.info("shruzard - ProxyBackendHandler: Unable to write back to App fabric... Closing channel "); + + future.channel().close(); + } + }); + } + + @Override + public void channelWritabilityChanged(ChannelHandlerContext ctx) { + // 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()); + } + 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); + } + + @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 new file mode 100644 index 000000000000..8a3c8fdcd649 --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/ProxyFrontendHandler.java @@ -0,0 +1,317 @@ +/* + * 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.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; + +/** + * 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 Channel outboundChannel; + private boolean connecting = false; + private boolean rejecting = 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) { + LOG.info("shruzard - ProxyFrontendHandler Received request!"); + HttpRequest req = (HttpRequest) msg; + + // 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) { + 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(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 + // Perform lock-free Compare-And-Swap evaluation using the AtomicReference loop + for (Map.Entry entry : podRegistry.entrySet()) { + PodState state = entry.getValue(); + + 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 (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(); + + if (state.tryClaimIdleLease(targetNamespace, idleTimeoutNanos)) { + targetWorkerAddress = entry.getKey(); + LOG.info("shruzard - ProxyFrontendHandler: Claimed idle pod " + + "for new namespace '{}' at {}", + targetNamespace, targetWorkerAddress); + break; + } + } + } + + // 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); + rejecting = true; + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, HttpResponseStatus.TOO_MANY_REQUESTS); + response.headers().set("Content-Length", "0"); + response.headers().set("Connection", "close"); + ctx.writeAndFlush(response); + ReferenceCountUtil.release(msg); + return; + } + + 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. + 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) + .option(ChannelOption.SO_KEEPALIVE, true) + .handler(new ChannelInitializer() { + @Override + 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 + LOG.info("shruzard - ProxyFrontendHandler: Connected with task worker successfully! "); + 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); + Object pendingMsg = pendingMessages.poll(); + while (pendingMsg != null) { + ReferenceCountUtil.release(pendingMsg); + pendingMsg = pendingMessages.poll(); + } + // Decrement inflight count on failed connection + PodState fallbackState = podRegistry.get(chosenWorker); + if (fallbackState != null) { + fallbackState.decrementInflightRequests(); + } + ctx.channel().close(); + } + }); + + // 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) { + ctx.channel().close(); + } + 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); + } + } + } + + @Override + public void channelReadComplete(ChannelHandlerContext ctx) { + // Flush any buffered outbound data to the worker socket + if (outboundChannel != null && outboundChannel.isActive() && !connecting) { + outboundChannel.flush(); + } + ctx.fireChannelReadComplete(); + } + + @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()); + } + ctx.fireChannelWritabilityChanged(); + } + + @Override + public void channelInactive(ChannelHandlerContext ctx) { + // When client closes connection, cleanly close the outbound worker socket + if (outboundChannel != null) { + closeOnFlush(outboundChannel); + } + } + + @Override + public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { + cause.printStackTrace(); + 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/RemoteClient.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/RemoteClient.java index af8a37222fed..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 @@ -20,10 +20,12 @@ 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; 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; @@ -45,13 +47,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 +66,10 @@ 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 Gson GSON = new Gson(); private final InternalAuthenticator internalAuthenticator; private final EndpointStrategy endpointStrategy; @@ -66,18 +77,37 @@ public class RemoteClient { private final String discoverableServiceName; 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; + 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; + } /** @@ -93,6 +123,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) { @@ -204,6 +242,7 @@ 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); if (httpResponse.getResponseCode() != HttpURLConnection.HTTP_OK) { @@ -260,16 +299,34 @@ public HttpURLConnection openConnection(HttpMethod method, String resource) thro * @throws ServiceUnavailableException if the service could not be discovered */ public URL resolve(String resource) { + 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) { 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); + } URI uri = URIScheme.createURI(discoverable, "%s%s", basePath, resource); try { + 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) { - // 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/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 bb1f11fd1f34..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 @@ -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; @@ -50,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; @@ -65,22 +68,30 @@ */ 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"; 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 RemoteClient fallbackClient; private final RetryStrategy retryStrategy; private final Predicate retryablePredicate; private final MetricsCollectionService metricsCollectionService; private final AeadCipher userEncryptionAeadCipher; 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) { @@ -92,13 +103,27 @@ public RemoteTaskExecutor(CConfiguration cConf, MetricsCollectionService metrics RemoteClientFactory remoteClientFactory, Type workerType, 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; + String taskServiceName = cConf.getBoolean(Constants.Security.Authorization.ENABLED) + ? Constants.Service.TASK_MANAGER : Constants.Service.TASK_WORKER; + this.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); + 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( + 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; if (workerType == Type.TASK_WORKER) { this.workerUrl = TASK_WORKER_URL; this.retryStrategy = RetryStrategies.fromConfiguration(cConf, @@ -126,12 +151,46 @@ 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) -> { try { - HttpRequest.Builder requestBuilder = remoteClient - .requestBuilder(HttpMethod.POST, workerUrl) + // 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) { + String embeddedNamespace = runnableTaskRequest.getParam().getEmbeddedTaskRequest().getNamespace(); + if (embeddedNamespace != null && !embeddedNamespace.isEmpty()) { + namespace = embeddedNamespace; + LOG.info("shruzard - RemoteTaskExecutor: Mapped SystemAppTask namespace to embedded: {}", + namespace); + } + } + 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; // 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. + 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()); if (compression) { requestBuilder.addHeader(HttpHeaders.CONTENT_ENCODING, "gzip"); @@ -148,16 +207,32 @@ public byte[] runTask(RunnableTaskRequest runnableTaskRequest) throws Exception } HttpRequest httpRequest = requestBuilder.build(); - HttpResponse httpResponse = remoteClient.execute(httpRequest); + + 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 + // 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("Received response code %s for %s", httpResponse.getResponseCode(), + 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) { @@ -168,11 +243,23 @@ 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( 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) { @@ -183,6 +270,14 @@ 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( + 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/StickyLeaseManager.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java new file mode 100644 index 000000000000..93a17f5ea37d --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/StickyLeaseManager.java @@ -0,0 +1,227 @@ +/* + * 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 java.util.function.Consumer; +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; + private final Consumer onLeaseAcquired; + private final Runnable onLeaseReleased; + + public StickyLeaseManager() { + this(10, 10, null, null); + } + + 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; + } + + /** + * 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( + "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; + } + } + + // 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; + } + + // 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; + } + + /** + * 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; + } + + + 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(); + totalTasksProcessedInLease.incrementAndGet(); + } + } + + /** + * 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 void enforceInactivityReclamation() { + NamespaceId leased = currentLease.get(); + if (leased != null && activeTaskCount.get() == 0) { + long idleDurationMillis = System.currentTimeMillis() - lastActivityTimeMillis; + + if (idleDurationMillis >= 600000L) { // 10 minutes + releaseLease(String.format("Security boundary hard-timeout (10 minutes) " + + "exceeded for %s", leased.getNamespace())); + } + } + } + + /** + * 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("shruzard - StickyLeaseManager: Release Lease (Logical Reset): " + + "Cleared namespace context for '{}'. Reason: {}", + oldNamespace.getNamespace(), reason); + if (onLeaseReleased != null) { + onLeaseReleased.run(); + } + } + } + + @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/TaskManagerService.java b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java new file mode 100644 index 000000000000..fe87a172340c --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerService.java @@ -0,0 +1,159 @@ +/* + * 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.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; +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; + +/** + * 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 { + + private static final Logger LOG = LoggerFactory.getLogger(TaskManagerService.class); + + private final int port; + private final String address; + private EventLoopGroup bossGroup; + 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; + private final DiscoveryService discoveryService; + private Cancellable cancellable; + + @Inject + 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; + + LOG.info("shruzard - Initializing TaskManagerService (Netty Proxy POC) on {}:{}", address, port); + } + + @Override + 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()); + 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(); + + // 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(); + } + 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 new file mode 100644 index 000000000000..44e567e8041b --- /dev/null +++ b/cdap-common/src/main/java/io/cdap/cdap/common/internal/remote/TaskManagerServiceModule.java @@ -0,0 +1,32 @@ +/* + * 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 that binds {@link TaskManagerService} in Singleton scope for dependency injection. + */ +public class TaskManagerServiceModule extends AbstractModule { + + @Override + protected void configure() { + // Bind the Netty Proxy service as a singleton + 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 c2b2b52757cf..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,11 +81,13 @@ public class TaskWorkerHttpHandlerInternal extends AbstractHttpHandler { BasicThrowable.class, new BasicThrowableCodec()).create(); private final RunnableTaskLauncher runnableTaskLauncher; + private final ScheduledExecutorService leaseReclamationExecutor; private final BiConsumer taskCompletionConsumer; /** * 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 +101,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 +120,34 @@ 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, + (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); + } + } + ); + + this.leaseReclamationExecutor = Executors.newSingleThreadScheduledExecutor( + Threads.createDaemonThreadFactory("lease-reclamation")); + this.leaseReclamationExecutor.scheduleAtFixedRate( + this.stickyLeaseManager::enforceInactivityReclamation, + 1, 1, TimeUnit.MINUTES); + + // Restart the service to clean up and re-claim resources after user code // execution. @@ -131,6 +156,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 +184,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 +207,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 +244,114 @@ private void stopAndShutdown(ScheduledExecutorService executorService, Consumer< @POST @Path("/run") public void run(FullHttpRequest request, HttpResponder responder) { + LOG.info("shruzard - 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); + + 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 { + RunnableTaskContext runnableTaskContext = new RunnableTaskContext(runnableTaskRequest); + runnableTaskLauncher.launchRunnableTask(runnableTaskContext); + + TaskDetails taskDetails = new TaskDetails(metricsCollectionService, + 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) + .add("X-Active-Tasks", String.valueOf(stickyLeaseManager.getActiveTaskCount())) + .add("X-Leased-Namespace", currentLease)); + } 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)); + } + } + + 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 +385,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 +394,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 +419,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 { 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..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 @@ -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; @@ -86,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; @@ -171,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) { @@ -194,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(); + } } /** @@ -224,7 +275,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 +593,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 +601,19 @@ Set toDiscoverables(String name, V1Service service, name); return Collections.emptySet(); } - hostname = ipAddr.get(); - } else { - hostname = String.format("%s.%s", meta.getName(), namespace); - } - - // We don't expect there is more than one service port, hence only pick the first one + String hostname = ipAddr.get(); + return servicePorts.stream() + .map(port -> createDiscoverable( + name, hostname, + port, payload) + ) + .filter(Objects::nonNull) + .findFirst() + .map(Collections::singleton) + .orElse(Collections.emptySet()); + } + + String hostname = String.format("%s.%s", meta.getName(), namespace); return servicePorts.stream() .map(port -> createDiscoverable( name, hostname, @@ -656,4 +713,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)); + } + } } 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/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerServiceMain.java b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerServiceMain.java new file mode 100644 index 000000000000..b8d3bd7b177c --- /dev/null +++ b/cdap-master/src/main/java/io/cdap/cdap/master/environment/k8s/TaskManagerServiceMain.java @@ -0,0 +1,83 @@ +/* + * 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 (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 TaskManagerServiceMain extends AbstractServiceMain { + + public static void main(String[] args) throws Exception { + main(TaskManagerServiceMain.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/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 new file mode 100644 index 000000000000..a1036c83a1ac --- /dev/null +++ b/task-manager-service.yaml @@ -0,0 +1,125 @@ +# 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: + name: cdap-task-manager + namespace: default + labels: + cdap.service: task.manager + cdap.instance: shru-basic-rbac-01 +spec: + replicas: 1 + selector: + matchLabels: + cdap.service: task.manager + template: + metadata: + labels: + cdap.service: task.manager + cdap.instance: shru-basic-rbac-01 + spec: + 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 + imagePullPolicy: Always + args: + - "io.cdap.cdap.master.environment.k8s.TaskManagerServiceMain" + - "--env=k8s" + env: + - name: SERVICE_NAME + value: task-manager + - name: OPTS + value: "-Xmx1024m -XX:MaxDirectMemorySize=768m -agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=5005" + ports: + - containerPort: 11025 + name: http + - containerPort: 5005 + name: debug + resources: + requests: + cpu: "500m" + memory: "1024Mi" + 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 + - mountPath: /cdap_configmap + name: cdap-cm-vol-cdap-shru-basic-rbac-01-configmap + 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-shru-basic-rbac-01-cconf + name: cdap-conf + - configMap: + defaultMode: 420 + name: cdap-shru-basic-rbac-01-hconf + name: hadoop-conf + - name: cdap-security + secret: + defaultMode: 420 + secretName: cdap-security + - configMap: + defaultMode: 420 + name: cdap-shru-basic-rbac-01-configmap + name: cdap-cm-vol-cdap-shru-basic-rbac-01-configmap +--- +apiVersion: v1 +kind: Service +metadata: + name: cdap-task-manager + namespace: default +spec: + selector: + cdap.service: task.manager + ports: + - protocol: TCP + port: 11025 + targetPort: 11025 + name: http + - protocol: TCP + port: 5005 + targetPort: 5005 + name: debug + type: ClusterIP