fix: Restarted the watch after stream reset #19055

extensions-core/kubernetes-extensions/pom.xml

@@ -111,6 +111,12 @@
       <artifactId>easymock</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.squareup.okhttp3</groupId>
+      <artifactId>mockwebserver3</artifactId>
+      <version>${okhttp.version}</version>
+      <scope>test</scope>
+    </dependency>
 
     <!-- others -->
     <dependency>

extensions-core/kubernetes-extensions/src/main/java/org/apache/druid/k8s/discovery/ChannelResetException.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.k8s.discovery;
+
+import java.io.IOException;
+
+public class ChannelResetException extends IOException
+{
+  public ChannelResetException(Throwable ex)
+  {
+    super(ex);
+  }
+}

extensions-core/kubernetes-extensions/src/main/java/org/apache/druid/k8s/discovery/DefaultK8sApiClient.java

@@ -40,6 +40,7 @@
 import org.apache.druid.java.util.common.logger.Logger;
 
 import java.io.IOException;
+import java.io.InterruptedIOException;
 import java.net.SocketTimeoutException;
 import java.util.HashMap;
 import java.util.List;
@@ -194,7 +195,7 @@
         private Watch.Response<DiscoveryDruidNodeAndResourceVersion> obj;
 
         @Override
-        public boolean hasNext() throws SocketTimeoutException
+        public boolean hasNext() throws IOException
         {
           try {
             while (watch.hasNext()) {
@@ -263,8 +264,14 @@
                 return true;
               } else if (item != null && item.type != null && item.type.equals(WatchResult.BOOKMARK)) {
                 // Events with type BOOKMARK will only contain resourceVersion and no metadata. See
-                // Kubernetes API documentation for details.
-                LOGGER.debug("BOOKMARK event fired, no nothing, only update resourceVersion");
+                // Kubernetes API documentation for details. Build a response carrying just the
+                //updated resourceVersion; otherwise obj would still hold the previous event (or
+                // null on the first event), corrupting the caller's resourceVersion tracking.
+                LOGGER.debug("BOOKMARK event fired, only updating resourceVersion");
+                obj = new Watch.Response<>(
+                        WatchResult.BOOKMARK,
+                        new DiscoveryDruidNodeAndResourceVersion(
+                                item.object.getMetadata().getResourceVersion(), null));
                 return true;
               } else {
                 LOGGER.error("WTH! item or item.type is NULL");
@@ -274,9 +281,11 @@
           catch (RuntimeException ex) {
             if (ex.getCause() instanceof SocketTimeoutException) {
               throw (SocketTimeoutException) ex.getCause();
-            } else {
-              throw ex;
             }
+            if (ex.getCause() instanceof IOException && !(ex.getCause() instanceof InterruptedIOException)) {
+              throw new ChannelResetException(ex.getCause());
+            }
+            throw ex;
           }
 
           return false;

extensions-core/kubernetes-extensions/src/main/java/org/apache/druid/k8s/discovery/K8sDruidNodeDiscoveryProvider.java

@@ -254,50 +254,51 @@ private void watch()
     private void keepWatching(String labelSelector, String resourceVersion)
     {
       String nextResourceVersion = resourceVersion;
-      while (lifecycleLock.isStarted()) {
-        try {
-          WatchResult iter =
-              k8sApiClient.watchPods(podInfo.getPodNamespace(), labelSelector, nextResourceVersion, nodeRole);
 
+      while (lifecycleLock.isStarted()) {
+        try (WatchResult iter = k8sApiClient.watchPods(podInfo.getPodNamespace(), labelSelector, nextResourceVersion, nodeRole)) {
           if (iter == null) {
             // history not available, we need to start from scratch
             return;
           }
 
-          try {
-            while (iter.hasNext()) {
-              Watch.Response<DiscoveryDruidNodeAndResourceVersion> item = iter.next();
-              if (item != null && item.type != null && item.object != null) {
-                switch (item.type) {
-                  case WatchResult.ADDED:
-                    baseNodeRoleWatcher.childAdded(item.object.getNode());
-                    break;
-                  case WatchResult.DELETED:
-                  case WatchResult.NOT_READY:
-                    // Use skipIfUnknown=true for all k8s discovery removals.
-                    // DELETED can fire after NOT_READY (so the service is already removed), or before ADDED (pod deleted before becoming ready).
-                    // NOT_READY can repeat during CrashLoopBackOff. None of these warrant the error-level logging that
-                    // comes with trying to remove an unknown service.
-                    baseNodeRoleWatcher.childRemoved(item.object.getNode(), true);
-                    break;
-                  default:
-                }
-
-                // This should be updated after the action has been dealt with successfully
-                nextResourceVersion = item.object.getResourceVersion();
-
-              } else {
-                // Try again by starting the watch from the beginning. This can happen if the
-                // watch goes bad.
-                LOGGER.debug("Received NULL item while watching role[%s]. Restarting watch.", this.nodeRole);
-                return;
-              }
+          while (iter.hasNext()) {
+            Watch.Response<DiscoveryDruidNodeAndResourceVersion> item = iter.next();
+
+            if (item == null || item.type == null || item.object == null) {
+              LOGGER.debug("Received NULL item while watching role[%s]. Restarting watch.", this.nodeRole);
+              return;
             }
-          }
-          finally {
-            iter.close();
+
+            switch (item.type) {
+              case WatchResult.ADDED:
+                baseNodeRoleWatcher.childAdded(item.object.getNode());
+                break;
+              case WatchResult.DELETED:
+              case WatchResult.NOT_READY:
+                // Use skipIfUnknown=true for all k8s discovery removals.
+                // DELETED can fire after NOT_READY (so the service is already removed), or before ADDED (pod deleted before becoming ready).
+                // NOT_READY can repeat during CrashLoopBackOff. None of these warrant the error-level logging that
+                // comes with trying to remove an unknown service.
+                baseNodeRoleWatcher.childRemoved(item.object.getNode(), true);
+                break;
+              default:
+            }
+
+            // This should be updated after the action has been dealt with successfully
+            nextResourceVersion = item.object.getResourceVersion();
           }
 
+          LOGGER.trace("Watch closed normally for role[%s]", this.nodeRole);
+          return;
+        }
+        catch (ChannelResetException ex) {
+          LOGGER.debug("Watch stream terminated normally for role[%s], restarting", this.nodeRole);
+          // Apply the same backoff as other retry paths below: a persistently reset stream
+          // (flaky proxy, LB idle timeout shorter than expected) would otherwise cause an
+          // unbounded relist/watch loop against the API server.
+          sleep(watcherErrorRetryWaitMS);
+          return;
         }
         catch (SocketTimeoutException ex) {
           // socket read timeout can happen normally due to k8s not having anything new to push leading to socket

extensions-core/kubernetes-extensions/src/main/java/org/apache/druid/k8s/discovery/WatchResult.java

@@ -21,15 +21,15 @@
 
 import io.kubernetes.client.util.Watch;
 
-import java.net.SocketTimeoutException;
+import java.io.IOException;
 
 /**
  * Iterator over k8s pod watch events that is aligned with the needs of Druid service discovery rather than
  * raw Kubernetes watch semantics. Implementations may remap or synthesize event types: for example,
  * a Kubernetes MODIFIED event for a pod whose containers are no longer ready (according to k8s readiness state) is surfaced as the
  * synthetic {@link #NOT_READY} type so that consumers can handle it as a removal from k8s service discovery.
  */
-public interface WatchResult
+public interface WatchResult extends AutoCloseable
 {
   String ADDED = "ADDED";
   String MODIFIED = "MODIFIED";
@@ -44,9 +44,7 @@ public interface WatchResult
    */
   String NOT_READY = "NOT_READY";
 
-  boolean hasNext() throws SocketTimeoutException;
+  boolean hasNext() throws IOException;
 
   Watch.Response<DiscoveryDruidNodeAndResourceVersion> next();
-
-  void close();
 }

extensions-core/kubernetes-extensions/src/test/java/org/apache/druid/k8s/discovery/DefaultK8sApiClientWatchPodsTest.java

@@ -0,0 +1,164 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.k8s.discovery;
+
+import com.fasterxml.jackson.databind.InjectableValues;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.kubernetes.client.openapi.ApiClient;
+import io.kubernetes.client.util.Watch;
+import mockwebserver3.MockResponse;
+import mockwebserver3.MockWebServer;
+import mockwebserver3.SocketEffect;
+import okhttp3.Headers;
+import org.apache.druid.discovery.DiscoveryDruidNode;
+import org.apache.druid.discovery.NodeRole;
+import org.apache.druid.jackson.DefaultObjectMapper;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.server.DruidNode;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+
+/**
+ * Exercises {@link DefaultK8sApiClient#watchPods} end to end against a real (but local, canned)
+ * HTTP watch stream: BOOKMARK events must surface their resourceVersion to the caller instead of
+ * a stale or null response, and a mid-stream connection reset must surface as
+ * {@link ChannelResetException} rather than an unrelated exception type.
+ */
+public class DefaultK8sApiClientWatchPodsTest
+{
+  private static final ObjectMapper JSON_MAPPER = new DefaultObjectMapper();
+
+  static {
+    // DiscoveryDruidNode's @JsonCreator takes a @JacksonInject ObjectMapper, so the mapper
+    // used to deserialize it must know how to inject itself.
+    JSON_MAPPER.setInjectableValues(new InjectableValues.Std().addValue(ObjectMapper.class, JSON_MAPPER));
+  }
+
+  private final DiscoveryDruidNode testNode = new DiscoveryDruidNode(
+      new DruidNode("druid/router", "test-host", true, 80, null, true, false), NodeRole.ROUTER, null);
+
+  private MockWebServer server;
+  private DefaultK8sApiClient client;
+
+  @BeforeEach
+  public void setUp() throws IOException
+  {
+    server = new MockWebServer();
+    server.start();
+
+    ApiClient apiClient = new ApiClient().setBasePath(server.url("/").toString());
+    client = new DefaultK8sApiClient(apiClient, JSON_MAPPER);
+  }
+
+  @AfterEach
+  public void tearDown() throws IOException
+  {
+    server.close();
+  }
+
+  @Test
+  public void testBookmarkAsFirstEventSurfacesResourceVersion() throws Exception
+  {
+    enqueueWatchStream(bookmarkLine("100"));
+
+    try (WatchResult result = client.watchPods("testns", null, "99", NodeRole.ROUTER)) {
+      Assertions.assertTrue(result.hasNext());
+      Watch.Response<DiscoveryDruidNodeAndResourceVersion> response = result.next();
+
+      Assertions.assertNotNull(response, "BOOKMARK event must not surface as a null response");
+      Assertions.assertEquals(WatchResult.BOOKMARK, response.type);
+      Assertions.assertNotNull(response.object);
+      Assertions.assertEquals("100", response.object.getResourceVersion());
+    }
+  }
+
+  @Test
+  public void testBookmarkAfterRealEventSurfacesItsOwnResourceVersion() throws Exception
+  {
+    enqueueWatchStream(addedLine("200", testNode) + bookmarkLine("201"));
+
+    try (WatchResult result = client.watchPods("testns", null, "99", NodeRole.ROUTER)) {
+      Assertions.assertTrue(result.hasNext());
+      Watch.Response<DiscoveryDruidNodeAndResourceVersion> added = result.next();
+      Assertions.assertEquals(WatchResult.ADDED, added.type);
+      Assertions.assertEquals("200", added.object.getResourceVersion());
+
+      Assertions.assertTrue(result.hasNext());
+      Watch.Response<DiscoveryDruidNodeAndResourceVersion> bookmark = result.next();
+
+      Assertions.assertNotNull(bookmark, "BOOKMARK event must not surface as a null response");
+      Assertions.assertEquals(WatchResult.BOOKMARK, bookmark.type);
+      Assertions.assertNotNull(bookmark.object);
+      Assertions.assertEquals(
+          "201",
+          bookmark.object.getResourceVersion(),
+          "BOOKMARK must carry its own resourceVersion, not the preceding event's");
+    }
+  }
+
+  @Test
+  public void testConnectionResetMidStreamSurfacesAsChannelResetException() throws Exception
+  {
+    // Truncate the body mid-object and sever the connection while it's being read, simulating
+    // a server-side stream reset. okhttp surfaces this as a ProtocolException, which
+    // DefaultK8sApiClient must remap to ChannelResetException rather than an unrelated
+    // RuntimeException, so callers can distinguish it from other IO failures.
+    server.enqueue(new MockResponse.Builder()
+        .code(200)
+        .headers(new Headers.Builder()
+            .add("Content-Type", "application/json")
+            .build())
+        .body(addedLine("200", testNode).substring(0, 20))
+        .onResponseBody(new SocketEffect.CloseSocket(true, true, true))
+        .build());
+
+    try (WatchResult result = client.watchPods("testns", null, "99", NodeRole.ROUTER)) {
+      Assertions.assertThrows(ChannelResetException.class, result::hasNext);
+    }
+  }
+
+  private void enqueueWatchStream(String body)
+  {
+    server.enqueue(new MockResponse(
+        200,
+        new Headers.Builder().add("Content-Type", "application/json").build(),
+        body));
+  }
+
+  private static String bookmarkLine(String resourceVersion)
+  {
+    return "{\"type\":\"BOOKMARK\",\"object\":{\"kind\":\"Pod\",\"apiVersion\":\"v1\",\"metadata\":{\"resourceVersion\":\""
+        + resourceVersion + "\"}}}\n";
+  }
+
+  private static String addedLine(String resourceVersion, DiscoveryDruidNode node) throws Exception
+  {
+    String infoAnnotation = K8sDruidNodeAnnouncer.getInfoAnnotation(node.getNodeRole());
+    String nodeJson = StringUtils.replace(JSON_MAPPER.writeValueAsString(node), "\"", "\\\"");
+    return "{\"type\":\"ADDED\",\"object\":{\"kind\":\"Pod\",\"apiVersion\":\"v1\",\"metadata\":{"
+        + "\"name\":\"test-pod\",\"resourceVersion\":\"" + resourceVersion + "\","
+        + "\"annotations\":{\"" + infoAnnotation + "\":\"" + nodeJson + "\"}},"
+        + "\"status\":{\"containerStatuses\":[{\"name\":\"main\",\"ready\":true}]}}}\n";
+  }
+}