Merge pull request #1118 from yue9944882/custom-work-queue-key-func

Support custom workqueue-keyfunc to directly manipulate work-queue instance

Author: k8s-ci-robot

extended/src/main/java/io/kubernetes/client/extended/controller/DefaultController.java

@@ -199,6 +199,15 @@ private void worker() {
     }
   }
 
+  public RateLimitingQueue<Request> getWorkQueue() {
+    return workQueue;
+  }
+
+  public DefaultController setWorkQueue(RateLimitingQueue<Request> workQueue) {
+    this.workQueue = workQueue;
+    return this;
+  }
+
   public String getName() {
     return name;
   }

extended/src/main/java/io/kubernetes/client/extended/controller/DefaultControllerWatch.java

@@ -94,22 +94,31 @@ public ResourceEventHandler<ApiType> getResourceEventHandler() {
       @Override
       public void onAdd(ApiType obj) {
         if (onAddFilterPredicate == null || onAddFilterPredicate.test(obj)) {
-          workQueue.add(workKeyGenerator.apply(obj));
+          Request req = workKeyGenerator.apply(obj);
+          if (null != req) {
+            workQueue.add(req);
+          }
         }
       }
 
       @Override
       public void onUpdate(ApiType oldObj, ApiType newObj) {
         if (onUpdateFilterPredicate == null || onUpdateFilterPredicate.test(oldObj, newObj)) {
-          workQueue.add(workKeyGenerator.apply(newObj));
+          Request req = workKeyGenerator.apply(newObj);
+          if (null != req) {
+            workQueue.add(req);
+          }
         }
       }
 
       @Override
       public void onDelete(ApiType obj, boolean deletedFinalStateUnknown) {
         if (onDeleteFilterPredicate == null
             || onDeleteFilterPredicate.test(obj, deletedFinalStateUnknown)) {
-          workQueue.add(workKeyGenerator.apply(obj));
+          Request req = workKeyGenerator.apply(obj);
+          if (null != req) {
+            workQueue.add(req);
+          }
         }
       }
     };

spring/src/main/java/io/kubernetes/client/spring/extended/controller/KubernetesReconcilerProcessor.java

@@ -12,6 +12,7 @@
 */
 package io.kubernetes.client.spring.extended.controller;
 
+import io.kubernetes.client.common.KubernetesObject;
 import io.kubernetes.client.extended.controller.Controller;
 import io.kubernetes.client.extended.controller.ControllerManager;
 import io.kubernetes.client.extended.controller.builder.ControllerBuilder;
@@ -20,6 +21,7 @@
 import io.kubernetes.client.extended.controller.reconciler.Request;
 import io.kubernetes.client.extended.workqueue.DefaultRateLimitingQueue;
 import io.kubernetes.client.extended.workqueue.RateLimitingQueue;
+import io.kubernetes.client.extended.workqueue.WorkQueue;
 import io.kubernetes.client.informer.SharedInformerFactory;
 import io.kubernetes.client.spring.extended.controller.annotation.*;
 import java.lang.reflect.InvocationTargetException;
@@ -99,15 +101,28 @@ private Controller buildController(SharedInformerFactory sharedInformerFactory,
     List<ReadyFuncAdaptor> readyFuncs = getReadyFuncs(r);
     for (KubernetesReconcilerWatch watch : watches.value()) {
       try {
-        Function<?, Request> workQueueKeyFunc = watch.workQueueKeyFunc().newInstance();
+
+        Function<? extends KubernetesObject, Request> workQueueKeyFunc;
+        try {
+          workQueueKeyFunc =
+              watch.workQueueKeyFunc().getConstructor(WorkQueue.class).newInstance(workQueue);
+        } catch (NoSuchMethodException e) {
+          workQueueKeyFunc = watch.workQueueKeyFunc().newInstance();
+        } catch (InvocationTargetException e) {
+          throw new BeanCreationException(
+              "Failed instantiating controller watch: " + e.getMessage());
+        }
+
+        final Function<? extends KubernetesObject, Request> finalWorkQueueKeyFunc =
+            workQueueKeyFunc;
         builder =
             builder.watch(
                 (q) -> {
                   return ControllerBuilder.controllerWatchBuilder(watch.apiTypeClass(), q)
                       .withOnAddFilter(addFilters.get(watch.apiTypeClass()))
                       .withOnUpdateFilter(updateFilters.get(watch.apiTypeClass()))
                       .withOnDeleteFilter(deleteFilters.get(watch.apiTypeClass()))
-                      .withWorkQueueKeyFunc(workQueueKeyFunc)
+                      .withWorkQueueKeyFunc(finalWorkQueueKeyFunc)
                       .withResyncPeriod(Duration.ofMillis(watch.resyncPeriodMillis()))
                       .build();
                 });

spring/src/main/java/io/kubernetes/client/spring/extended/controller/annotation/KubernetesReconcilerWatch.java

@@ -38,9 +38,14 @@
    * Work queue key func class maps the source resource of the watch event to a standard reconciler
    * request.
    *
+   * <p>Optionally you can declare the constructor of the class to receive a {@link
+   * io.kubernetes.client.extended.workqueue.WorkQueue} in the parameter in order to customize the
+   * work-queue key-func.
+   *
    * @return the class
    */
-  Class<? extends Function<?, Request>> workQueueKeyFunc() default DefaultReflectiveKeyFunc.class;
+  Class<? extends Function<? extends KubernetesObject, Request>> workQueueKeyFunc() default
+      DefaultReflectiveKeyFunc.class;
 
   /**
    * Resync period in milliseconds .

spring/src/test/java/io/kubernetes/client/spring/extended/controller/KubernetesReconcilerCreatorTest.java

@@ -12,24 +12,41 @@
 */
 package io.kubernetes.client.spring.extended.controller;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.fail;
 
 import com.github.tomakehurst.wiremock.junit.WireMockRule;
+import io.kubernetes.client.common.KubernetesObject;
 import io.kubernetes.client.extended.controller.Controller;
+import io.kubernetes.client.extended.controller.DefaultController;
 import io.kubernetes.client.extended.controller.reconciler.Reconciler;
 import io.kubernetes.client.extended.controller.reconciler.Request;
 import io.kubernetes.client.extended.controller.reconciler.Result;
+import io.kubernetes.client.extended.workqueue.WorkQueue;
 import io.kubernetes.client.informer.SharedInformer;
+import io.kubernetes.client.informer.SharedInformerFactory;
+import io.kubernetes.client.informer.cache.DeltaFIFO;
 import io.kubernetes.client.informer.cache.Lister;
+import io.kubernetes.client.informer.impl.DefaultSharedIndexInformer;
+import io.kubernetes.client.openapi.models.V1ConfigMap;
+import io.kubernetes.client.openapi.models.V1ObjectMeta;
 import io.kubernetes.client.openapi.models.V1Pod;
-import io.kubernetes.client.spring.extended.controller.annotation.*;
+import io.kubernetes.client.openapi.models.V1PodList;
 import io.kubernetes.client.spring.extended.controller.annotation.AddWatchEventFilter;
 import io.kubernetes.client.spring.extended.controller.annotation.DeleteWatchEventFilter;
+import io.kubernetes.client.spring.extended.controller.annotation.KubernetesReconciler;
+import io.kubernetes.client.spring.extended.controller.annotation.KubernetesReconcilerReadyFunc;
+import io.kubernetes.client.spring.extended.controller.annotation.KubernetesReconcilerWatch;
+import io.kubernetes.client.spring.extended.controller.annotation.KubernetesReconcilerWatches;
 import io.kubernetes.client.spring.extended.controller.annotation.UpdateWatchEventFilter;
+import java.util.LinkedList;
+import java.util.function.Function;
+import javax.annotation.Resource;
+import org.apache.commons.lang3.tuple.MutablePair;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.boot.test.context.TestConfiguration;
 import org.springframework.context.annotation.Bean;
@@ -47,8 +64,12 @@ public class KubernetesReconcilerCreatorTest {
   static class TestConfig {
 
     @Bean
-    public TestReconciler podReconciler() {
-      return new TestReconciler();
+    public TestReconciler podReconciler(
+        SharedInformer<V1Pod> podInformer,
+        Lister<V1Pod> podLister,
+        SharedInformer<V1ConfigMap> configMapInformer,
+        Lister<V1ConfigMap> configMapLister) {
+      return new TestReconciler(podInformer, podLister, configMapLister);
     }
   }
 
@@ -58,20 +79,32 @@ public TestReconciler podReconciler() {
           @KubernetesReconcilerWatches({
             @KubernetesReconcilerWatch(
                 apiTypeClass = V1Pod.class,
+                workQueueKeyFunc = CustomWorkQueueKeyFunc.class,
                 resyncPeriodMillis = 60 * 1000L // resync every 60s
                 )
           }))
   static class TestReconciler implements Reconciler {
 
-    private int observedPodCount;
+    private Request receivedRequest;
 
-    @Autowired private SharedInformer<V1Pod> podInformer;
+    public TestReconciler(
+        SharedInformer<V1Pod> podInformer,
+        Lister<V1Pod> podLister,
+        Lister<V1ConfigMap> configMapLister) {
+      this.podInformer = podInformer;
+      this.podLister = podLister;
+      this.configMapLister = configMapLister;
+    }
+
+    private final SharedInformer<V1Pod> podInformer;
 
-    @Autowired private Lister<V1Pod> podLister;
+    private final Lister<V1Pod> podLister;
+
+    private final Lister<V1ConfigMap> configMapLister;
 
     @Override
     public Result reconcile(Request request) {
-      observedPodCount = podLister.list().size();
+      receivedRequest = request;
       return new Result(false);
     }
 
@@ -96,10 +129,53 @@ private boolean podInformerCacheReady() {
     }
   }
 
-  @Autowired private Controller testController;
+  @Resource(name = "test-reconcile")
+  private Controller testController;
+
+  @Resource private TestReconciler testReconciler;
+
+  @Resource private SharedInformerFactory sharedInformerFactory;
+
+  public static class CustomWorkQueueKeyFunc implements Function<KubernetesObject, Request> {
+
+    public CustomWorkQueueKeyFunc(WorkQueue<Request> workQueue) {
+      this.workQueue = workQueue;
+    }
+
+    private final WorkQueue<Request> workQueue;
+
+    @Override
+    public Request apply(KubernetesObject item) {
+      workQueue.add(new Request("foo"));
+      return null;
+    }
+  }
 
   @Test
-  public void testSimplePodController() {
+  public void testSimplePodController() throws InterruptedException {
     assertNotNull(testController);
+    assertNotNull(testReconciler);
+
+    sharedInformerFactory.startAllRegisteredInformers();
+
+    ((DefaultSharedIndexInformer<V1Pod, V1PodList>) testReconciler.podInformer)
+        .handleDeltas(
+            new LinkedList<MutablePair<DeltaFIFO.DeltaType, KubernetesObject>>() {
+              {
+                add(
+                    new MutablePair<>(
+                        DeltaFIFO.DeltaType.Added,
+                        new V1Pod().metadata(new V1ObjectMeta().namespace("a").name("b"))));
+              }
+            });
+
+    Thread.sleep(500);
+
+    WorkQueue<Request> workQueue = ((DefaultController) testController).getWorkQueue();
+    assertEquals(1, workQueue.length());
+    if (workQueue.length() != 1) {
+      fail();
+    }
+    assertEquals("foo", workQueue.get().getName());
   }
 }

util/src/main/java/io/kubernetes/client/informer/impl/DefaultSharedIndexInformer.java

@@ -216,7 +216,7 @@ public boolean hasSynced() {
    *
    * @param deltas deltas
    */
-  private void handleDeltas(Deque<MutablePair<DeltaFIFO.DeltaType, KubernetesObject>> deltas) {
+  public void handleDeltas(Deque<MutablePair<DeltaFIFO.DeltaType, KubernetesObject>> deltas) {
     if (CollectionUtils.isEmpty(deltas)) {
       return;
     }