From f30d58ad312c7f664dde032256efe33d31f5acb6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Xavier=20L=C3=A9aut=C3=A9?= <xavier@metamarkets.com>
Date: Thu, 29 May 2014 15:41:12 -0700
Subject: [PATCH] filter out segment callbacks not related to given filter

---
 .../client/BatchServerInventoryView.java      | 24 +++++++++++-----
 .../client/SingleServerInventoryView.java     | 28 +++++++++++++------
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/server/src/main/java/io/druid/client/BatchServerInventoryView.java b/server/src/main/java/io/druid/client/BatchServerInventoryView.java
index ebb808b3d73..f8225d4e905 100644
--- a/server/src/main/java/io/druid/client/BatchServerInventoryView.java
+++ b/server/src/main/java/io/druid/client/BatchServerInventoryView.java
@@ -135,7 +135,7 @@ public class BatchServerInventoryView extends ServerInventoryView<Set<DataSegmen
 
   @Override
   public void registerSegmentCallback(
-      Executor exec, final SegmentCallback callback, Predicate<DataSegment> filter
+      final Executor exec, final SegmentCallback callback, final Predicate<DataSegment> filter
   )
   {
     segmentPredicates.put(callback, filter);
@@ -147,9 +147,14 @@ public class BatchServerInventoryView extends ServerInventoryView<Set<DataSegmen
               DruidServerMetadata server, DataSegment segment
           )
           {
-            final CallbackAction action = callback.segmentAdded(server, segment);
-            if (action.equals(CallbackAction.UNREGISTER)) {
-              segmentPredicates.remove(callback);
+            final CallbackAction action;
+            if(filter.apply(segment)) {
+              action = callback.segmentAdded(server, segment);
+              if (action.equals(CallbackAction.UNREGISTER)) {
+                segmentPredicates.remove(callback);
+              }
+            } else {
+              action = CallbackAction.CONTINUE;
             }
             return action;
           }
@@ -159,9 +164,14 @@ public class BatchServerInventoryView extends ServerInventoryView<Set<DataSegmen
               DruidServerMetadata server, DataSegment segment
           )
           {
-            final CallbackAction action = callback.segmentRemoved(server, segment);
-            if (action.equals(CallbackAction.UNREGISTER)) {
-              segmentPredicates.remove(callback);
+            final CallbackAction action;
+            if(filter.apply(segment)) {
+              action = callback.segmentRemoved(server, segment);
+              if (action.equals(CallbackAction.UNREGISTER)) {
+                segmentPredicates.remove(callback);
+              }
+            } else {
+              action = CallbackAction.CONTINUE;
             }
             return action;
           }
diff --git a/server/src/main/java/io/druid/client/SingleServerInventoryView.java b/server/src/main/java/io/druid/client/SingleServerInventoryView.java
index 084f24d7437..7dd1e2f7452 100644
--- a/server/src/main/java/io/druid/client/SingleServerInventoryView.java
+++ b/server/src/main/java/io/druid/client/SingleServerInventoryView.java
@@ -99,7 +99,7 @@ public class SingleServerInventoryView extends ServerInventoryView<DataSegment>
 
   @Override
   public void registerSegmentCallback(
-      Executor exec, final SegmentCallback callback, Predicate<DataSegment> filter
+      final Executor exec, final SegmentCallback callback, final Predicate<DataSegment> filter
   )
   {
     segmentPredicates.put(callback, filter);
@@ -111,9 +111,14 @@ public class SingleServerInventoryView extends ServerInventoryView<DataSegment>
               DruidServerMetadata server, DataSegment segment
           )
           {
-            final CallbackAction action = callback.segmentAdded(server, segment);
-            if (action.equals(CallbackAction.UNREGISTER)) {
-              segmentPredicates.remove(callback);
+            final CallbackAction action;
+            if(filter.apply(segment)) {
+              action = callback.segmentAdded(server, segment);
+              if (action.equals(CallbackAction.UNREGISTER)) {
+                segmentPredicates.remove(callback);
+              }
+            } else {
+              action = CallbackAction.CONTINUE;
             }
             return action;
           }
@@ -123,11 +128,18 @@ public class SingleServerInventoryView extends ServerInventoryView<DataSegment>
               DruidServerMetadata server, DataSegment segment
           )
           {
-            final CallbackAction action = callback.segmentRemoved(server, segment);
-            if (action.equals(CallbackAction.UNREGISTER)) {
-              segmentPredicates.remove(callback);
+            {
+              final CallbackAction action;
+              if(filter.apply(segment)) {
+                action = callback.segmentRemoved(server, segment);
+                if (action.equals(CallbackAction.UNREGISTER)) {
+                  segmentPredicates.remove(callback);
+                }
+              } else {
+                action = CallbackAction.CONTINUE;
+              }
+              return action;
             }
-            return action;
           }
         }
     );