This is an automated email from the ASF dual-hosted git repository.

jamesshao pushed a commit to branch upsert
in repository https://gitbox.apache.org/repos/asf/incubator-pinot.git

commit 9ac12f38bcb30eed228737a1ae85445eac8a4d18
Author: james Shao <[email protected]>
AuthorDate: Wed Oct 23 16:50:35 2019 -0700

    fix issue in retention management and refactor
    
    Summary:
    right now if we delete and recreate the table, the seq number of segments 
got reset to start from 0.
    this cause issue in newer segment are considered as older than the last 
segment we have seen.
    this diff would solve such issue
    
    we also refactor a bunch of stuff to base class so the kc retention
    manager and server retention manager are mostly using the same logics
    
    Reviewers: tingchen, csoman, bzzhang
    
    Reviewed By: bzzhang
    
    Differential Revision: https://code.uberinternal.com/D3528221
---
 pinot-grigio/pinot-grigio-common/pom.xml           |   5 +
 .../UpdateLogRetentionManagerImpl.java             |   2 +-
 .../UpdateLogTableRetentionManager.java            |  98 ++++++++++----
 .../UpdateLogTableRetentionManagerImpl.java        |  79 +----------
 .../grigio}/common/utils/IdealStateHelper.java     |   2 +-
 .../UpdateLogTableRetentionManagerImplTest.java    | 119 +++++++++++++++++
 .../grigio/common/utils/IdealStateHelperTest.java  |  74 +++++++++++
 pinot-grigio/pinot-grigio-coordinator/pom.xml      |   5 +
 .../KCUpdateLogRetentionManagerImpl.java           |   2 +-
 .../KCUpdateLogTableRetentionManagerImpl.java      |  58 ++------
 .../starter/KeyCoordinatorStarter.java             |   2 +-
 .../KCUpdateLogRetentionManagerImplTest.java       | 114 ++++++++++++++++
 .../KCUpdateLogTableRetentionManagerImplTest.java  | 148 +++++++++++++++++++++
 .../pinot/server/starter/ServerInstance.java       |   2 +-
 14 files changed, 560 insertions(+), 150 deletions(-)

diff --git a/pinot-grigio/pinot-grigio-common/pom.xml 
b/pinot-grigio/pinot-grigio-common/pom.xml
index 251cb5a..53adbbf 100644
--- a/pinot-grigio/pinot-grigio-common/pom.xml
+++ b/pinot-grigio/pinot-grigio-common/pom.xml
@@ -60,6 +60,11 @@
       <artifactId>guava</artifactId>
     </dependency>
     <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
       <groupId>org.testng</groupId>
       <artifactId>testng</artifactId>
       <scope>test</scope>
diff --git 
a/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogRetentionManagerImpl.java
 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogRetentionManagerImpl.java
index 6699a64..3852f99 100644
--- 
a/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogRetentionManagerImpl.java
+++ 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogRetentionManagerImpl.java
@@ -18,7 +18,7 @@
  */
 package org.apache.pinot.grigio.common.storageProvider.retentionManager;
 
-import org.apache.pinot.common.utils.IdealStateHelper;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
diff --git 
a/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManager.java
 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManager.java
index afb42bc..62215b8 100644
--- 
a/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManager.java
+++ 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManager.java
@@ -20,54 +20,99 @@ package 
org.apache.pinot.grigio.common.storageProvider.retentionManager;
 
 import com.google.common.collect.ImmutableMap;
 import org.apache.pinot.common.utils.LLCSegmentName;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * class to help decide whether to keep update log for a particular segment or 
not
  */
-public interface UpdateLogTableRetentionManager {
+public abstract class UpdateLogTableRetentionManager {
 
-  /**
-   * decide whether we should ingest the update log for a given segment name
-   * @param segmentName
-   * @return true if we should keep the update log for a particular segment, 
false otherwise
-   */
-  boolean shouldIngestForSegment(String segmentName);
+  protected static final Logger LOGGER = 
LoggerFactory.getLogger(UpdateLogTableRetentionManagerImpl.class);
+
+  protected String _tableName;
+  protected IdealStateHelper _idealStateHelper;
+  // use the concurrent hashmap as concurrent hashset
+  protected Map<Integer, LLCSegmentName> _partitionToLastSegment = new 
ConcurrentHashMap<>();
+  protected Map<String, String> _blacklistedSegments = new 
ConcurrentHashMap<>();
+  protected Map<String, Map<String, String>> _segmentsToInstanceMap;
+
+  public UpdateLogTableRetentionManager(IdealStateHelper idealStateHelper, 
String tableName) {
+    _idealStateHelper = idealStateHelper;
+    _tableName = tableName;
+  }
 
   /**
    * for external components to notify retention manager that physical data 
for this segment has been deleted
    * @param segmentName the name of the segment that we are deleting from 
local storage
    */
-  void notifySegmentDeletion(String segmentName);
+  public abstract void notifySegmentDeletion(String segmentName);
 
   /**
    * for external components to notify us that segments for this tables has 
been change and we should refresh it
    */
-  void notifySegmentsChange();
+  public abstract void notifySegmentsChange();
+
+  /**
+   * check if the the given segment is assigned to current server
+   */
+  protected abstract boolean isSegmentAssignedToCurrentServer(String 
segmentName);
 
   /**
-   * get the mapping between the partition to the last segment (based on 
timestamp)
-   * @param segments the list of segment we are comparing
-   * @return the mapping between the partition id and the last segment of that 
partition
+   * update state from helix
+   */
+  protected void updateStateFromHelix() {
+    long start = System.currentTimeMillis();
+    _segmentsToInstanceMap = 
ImmutableMap.copyOf(_idealStateHelper.getSegmentToInstanceMap(_tableName));
+    if (_segmentsToInstanceMap.size() == 0) {
+      LOGGER.error("failed to get any segment for the current table {}", 
_tableName);
+    }
+    Map<Integer, LLCSegmentName> partitionLastSeg = new HashMap<>();
+    for (String segmentStr: _segmentsToInstanceMap.keySet()) {
+      LLCSegmentName llcSegmentName = new LLCSegmentName(segmentStr);
+      int partition = llcSegmentName.getPartitionId();
+      if (!partitionLastSeg.containsKey(partition)) {
+        partitionLastSeg.put(partition, llcSegmentName);
+      } else if (partitionLastSeg.get(partition).getSequenceNumber() < 
llcSegmentName.getSequenceNumber()) {
+        partitionLastSeg.put(partition, llcSegmentName);
+      }
+    }
+    _partitionToLastSegment = ImmutableMap.copyOf(partitionLastSeg);
+    LOGGER.info("updated table {} state from helix in {} ms", _tableName, 
System.currentTimeMillis() - start);
+  }
+
+  /**
+   * decide whether we should ingest the update log for a given segment name
+   * @param segmentName
+   * @return true if we should keep the update log for a particular segment, 
false otherwise
    */
-  static Map<Integer, LLCSegmentName> getPartitionToLastSegment(Set<String> 
segments) {
-    Map<Integer, LLCSegmentName> partitionToLastSegment = new HashMap<>();
-    for (String segment : segments) {
-      LLCSegmentName llcSegmentName = new LLCSegmentName(segment);
+  public boolean shouldIngestForSegment(String segmentName) {
+    if (_segmentsToInstanceMap.containsKey(segmentName)) {
+      return isSegmentAssignedToCurrentServer(segmentName);
+    } else if (_blacklistedSegments.containsKey(segmentName)) {
+      return false;
+    } else {
+      LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName);
       int partition = llcSegmentName.getPartitionId();
-      if (!partitionToLastSegment.containsKey(partition)) {
-        partitionToLastSegment.put(partition, llcSegmentName);
-      } else {
-        LLCSegmentName currLastSegment = partitionToLastSegment.get(partition);
-        if (llcSegmentName.getCreationTimeStamp() > 
currLastSegment.getCreationTimeStamp()) {
-          partitionToLastSegment.put(partition, llcSegmentName);
+      if (!_partitionToLastSegment.containsKey(llcSegmentName.getPartitionId())
+          || compareSegment(llcSegmentName, 
_partitionToLastSegment.get(partition))) {
+        // assume our idealState is out of date
+        updateStateFromHelix();
+        if (_segmentsToInstanceMap.containsKey(segmentName) && 
isSegmentAssignedToCurrentServer(segmentName)) {
+          LOGGER.info("segment {} matched in ideal state after refresh", 
segmentName);
+          return true;
         }
       }
+      // we most probably got a segment that is from a deleted table or 
segment assigned to another table
+      // assume we don't do re-balance, we won't do refresh
+      _blacklistedSegments.put(segmentName, segmentName);
+      return false;
     }
-    return ImmutableMap.copyOf(partitionToLastSegment);
   }
 
   /**
@@ -76,9 +121,8 @@ public interface UpdateLogTableRetentionManager {
    * @param segment2
    * @return true if the segment1 is a "newer" segment
    */
-  static boolean compareSegment(LLCSegmentName segment1, LLCSegmentName 
segment2) {
-    return segment1.getSequenceNumber() > segment2.getSequenceNumber() &&
-        segment1.getCreationTimeStamp() > segment2.getCreationTimeStamp();
+  protected boolean compareSegment(LLCSegmentName segment1, LLCSegmentName 
segment2) {
+    return segment1.getCreationTimeStamp() > segment2.getCreationTimeStamp();
   }
 
 }
diff --git 
a/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManagerImpl.java
 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManagerImpl.java
index f145eb5..91c098a 100644
--- 
a/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManagerImpl.java
+++ 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManagerImpl.java
@@ -18,88 +18,20 @@
  */
 package org.apache.pinot.grigio.common.storageProvider.retentionManager;
 
-import com.google.common.collect.ImmutableMap;
-import org.apache.pinot.common.utils.IdealStateHelper;
-import org.apache.pinot.common.utils.LLCSegmentName;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
 
 import java.util.Arrays;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
-public class UpdateLogTableRetentionManagerImpl implements 
UpdateLogTableRetentionManager {
+public class UpdateLogTableRetentionManagerImpl extends 
UpdateLogTableRetentionManager {
 
-  private static final Logger LOGGER = 
LoggerFactory.getLogger(UpdateLogTableRetentionManagerImpl.class);
-
-  protected IdealStateHelper _idealStateHelper;
-  protected String _tableName;
-
-  private Map<Integer, LLCSegmentName> _partitionLastSeg = new 
ConcurrentHashMap<>();
-  private Map<String, Map<String, String>> _segmentsToInstanceMap;
-  // use the concurrent hashmap as concurrent hashset
-  private Map<String, String> _blacklistedSegmentNames = new 
ConcurrentHashMap<>();
-  private String _instanceId;
+  protected String _instanceId;
 
   public UpdateLogTableRetentionManagerImpl(IdealStateHelper idealStateHelper, 
String tableName, String instanceId) {
-    _idealStateHelper = idealStateHelper;
-    _tableName = tableName;
+    super(idealStateHelper, tableName);
     _instanceId = instanceId;
     updateStateFromHelix();
   }
 
-  private void updateStateFromHelix() {
-    long start = System.currentTimeMillis();
-    _segmentsToInstanceMap = 
ImmutableMap.copyOf(_idealStateHelper.getSegmentToInstanceMap(_tableName));
-    if (_segmentsToInstanceMap.size() == 0) {
-      LOGGER.error("failed to get any segment for the current table {}", 
_tableName);
-    }
-    Map<Integer, LLCSegmentName> partitionLastSeg = new HashMap<>();
-    for (String segmentStr: _segmentsToInstanceMap.keySet()) {
-      LLCSegmentName llcSegmentName = new LLCSegmentName(segmentStr);
-      int partition = llcSegmentName.getPartitionId();
-      if (!partitionLastSeg.containsKey(partition)) {
-        partitionLastSeg.put(partition, llcSegmentName);
-      } else if (partitionLastSeg.get(partition).getSequenceNumber() < 
llcSegmentName.getSequenceNumber()) {
-        partitionLastSeg.put(partition, llcSegmentName);
-      }
-    }
-    _partitionLastSeg = ImmutableMap.copyOf(partitionLastSeg);
-    LOGGER.info("updated table {} state from helix in {} ms", _tableName, 
System.currentTimeMillis() - start);
-  }
-
-  @Override
-  public synchronized boolean shouldIngestForSegment(String segmentName) {
-    if (_segmentsToInstanceMap.containsKey(segmentName)) {
-      return isSegmentAssignedToCurrentServer(segmentName);
-    } else if (_blacklistedSegmentNames.containsKey(segmentName)) {
-      return false;
-    } else {
-      LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName);
-      int partitionId = llcSegmentName.getPartitionId();
-      int seq = llcSegmentName.getSequenceNumber();
-      long ts = llcSegmentName.getCreationTimeStamp();
-      if (_partitionLastSeg.containsKey(partitionId)) {
-        LLCSegmentName latestSegment = _partitionLastSeg.get(partitionId);
-        if (seq > latestSegment.getSequenceNumber() && ts > 
latestSegment.getCreationTimeStamp()) {
-          // assume our idealState is out of date
-          updateStateFromHelix();
-          if (isSegmentAssignedToCurrentServer(segmentName)) {
-            LOGGER.info("segment {} matched in ideal state after refresh", 
segmentName);
-            return true;
-          }
-        } else {
-          // we most probably got a segment that is from a deleted table or 
segment assigned to another table
-          // assume we don't do re-balance, we won't do refresh
-          LOGGER.info("adding segment {} to blacklist segment list", 
segmentName);
-          _blacklistedSegmentNames.put(segmentName, segmentName);
-        }
-      }
-    }
-    return false;
-  }
-
   @Override
   public synchronized void notifySegmentDeletion(String segmentName) {
     LOGGER.info("handling segment deletion in retention manager");
@@ -115,7 +47,8 @@ public class UpdateLogTableRetentionManagerImpl implements 
UpdateLogTableRetenti
     updateStateFromHelix();
   }
 
-  private boolean isSegmentAssignedToCurrentServer(String segmentName) {
+  @Override
+  protected boolean isSegmentAssignedToCurrentServer(String segmentName) {
     return _segmentsToInstanceMap.containsKey(segmentName)
         && _segmentsToInstanceMap.get(segmentName).containsKey(_instanceId)
         && 
!"DROPPED".equals(_segmentsToInstanceMap.get(segmentName).get(_instanceId));
diff --git 
a/pinot-common/src/main/java/org/apache/pinot/common/utils/IdealStateHelper.java
 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/utils/IdealStateHelper.java
similarity index 98%
rename from 
pinot-common/src/main/java/org/apache/pinot/common/utils/IdealStateHelper.java
rename to 
pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/utils/IdealStateHelper.java
index 9ebc72d..b0c193f 100644
--- 
a/pinot-common/src/main/java/org/apache/pinot/common/utils/IdealStateHelper.java
+++ 
b/pinot-grigio/pinot-grigio-common/src/main/java/org/apache/pinot/grigio/common/utils/IdealStateHelper.java
@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-package org.apache.pinot.common.utils;
+package org.apache.pinot.grigio.common.utils;
 
 import org.apache.helix.HelixAdmin;
 import org.apache.helix.model.IdealState;
diff --git 
a/pinot-grigio/pinot-grigio-common/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManagerImplTest.java
 
b/pinot-grigio/pinot-grigio-common/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManagerImplTest.java
new file mode 100644
index 0000000..83fba68
--- /dev/null
+++ 
b/pinot-grigio/pinot-grigio-common/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/UpdateLogTableRetentionManagerImplTest.java
@@ -0,0 +1,119 @@
+/**
+ * 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.pinot.grigio.common.storageProvider.retentionManager;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.mockito.internal.verification.VerificationModeFactory.times;
+
+public class UpdateLogTableRetentionManagerImplTest {
+
+  private IdealStateHelper mockIdealStateHelper;
+  private Map<String, Map<String, String>> segmentsInstanceMap;
+  private UpdateLogTableRetentionManagerImpl updateLogTableRetentionManager;
+
+  @BeforeMethod
+  public void init() {
+    segmentsInstanceMap = new HashMap<>();
+
+    segmentsInstanceMap.put("table__0__10__20191027T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server2", "ONLINE"));
+    segmentsInstanceMap.put("table__0__11__20191028T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server3", "ONLINE"));
+    segmentsInstanceMap.put("table__1__10__20191027T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server2", "ONLINE"));
+    segmentsInstanceMap.put("table__2__10__20191027T2041Z",
+        ImmutableMap.of("server2", "ONLINE", "server3", "ONLINE"));
+
+    mockIdealStateHelper = mock(IdealStateHelper.class);
+    
when(mockIdealStateHelper.getSegmentToInstanceMap("table")).thenReturn(segmentsInstanceMap);
+
+    updateLogTableRetentionManager = new 
UpdateLogTableRetentionManagerImpl(mockIdealStateHelper, "table",
+        "server1");
+  }
+
+  @Test
+  public void testShouldIngestForSegment() {
+    // test the happy paths
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__0__10__20191027T2041Z"));
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__1__10__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    // test for segments that has lower seq
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__1__5__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    //test for segment that are older and not in existing segment list
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__1__20__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    // test for partition that has no assignment on host
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__2__10__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    //test for older timestamp but higher seq
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__0__20__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    // test for newer segments, and there is update in newer ideal state
+    segmentsInstanceMap.put("table__0__12__20191029T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server3", "ONLINE"));
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__0__12__20191029T2041Z"));
+    verify(mockIdealStateHelper, times(2)).getSegmentToInstanceMap("table");
+
+    // test for newer segments and there is no update in newer ideal state
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__0__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(3)).getSegmentToInstanceMap("table");
+
+    // multiple attempts trying to fetch for segment should not recheck the 
ideal state often
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__0__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(3)).getSegmentToInstanceMap("table");
+
+    // for unknown partition, we should also refresh data
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__2__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(4)).getSegmentToInstanceMap("table");
+
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__3__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(5)).getSegmentToInstanceMap("table");
+  }
+
+  @Test
+  public void testNotifySegmentDeletion() {
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+    
updateLogTableRetentionManager.notifySegmentDeletion("table__0__10__20191027T2041Z");
+    verify(mockIdealStateHelper, times(2)).getSegmentToInstanceMap("table");
+  }
+
+  @Test
+  public void testNotifySegmentsChange() {
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+    
updateLogTableRetentionManager.notifySegmentDeletion("table__0__10__20191027T2041Z");
+    verify(mockIdealStateHelper, times(2)).getSegmentToInstanceMap("table");
+  }
+}
\ No newline at end of file
diff --git 
a/pinot-grigio/pinot-grigio-common/src/test/java/org/apache/pinot/grigio/common/utils/IdealStateHelperTest.java
 
b/pinot-grigio/pinot-grigio-common/src/test/java/org/apache/pinot/grigio/common/utils/IdealStateHelperTest.java
new file mode 100644
index 0000000..4245e5b
--- /dev/null
+++ 
b/pinot-grigio/pinot-grigio-common/src/test/java/org/apache/pinot/grigio/common/utils/IdealStateHelperTest.java
@@ -0,0 +1,74 @@
+/**
+ * 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.pinot.grigio.common.utils;
+
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.model.IdealState;
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+import java.util.Set;
+
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class IdealStateHelperTest {
+
+  private HelixAdmin mockAmin;
+  private IdealStateHelper idealStateHelper;
+  private IdealState mockIdealState;
+
+  @BeforeClass
+  public void init() {
+    mockIdealState = new IdealState("resource");
+    mockIdealState.enable(true);
+    mockIdealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
+    mockIdealState.setPartitionState("seg1", "host1", "ONLINE");
+    mockIdealState.setPartitionState("seg1", "host2", "ONLINE");
+    mockIdealState.setPartitionState("seg2", "host2", "ONLINE");
+    mockIdealState.setPartitionState("seg2", "host3", "OFFLINE");
+
+    mockAmin = mock(HelixAdmin.class);
+    when(mockAmin.getResourceIdealState(anyString(), 
anyString())).thenReturn(mockIdealState);
+    idealStateHelper = new IdealStateHelper(mockAmin, "clusterName");
+  }
+
+  @Test
+  public void testGetAllSegmentsForTable() {
+    Set<String> segments = idealStateHelper.getAllSegmentsForTable("table");
+    Assert.assertEquals(segments.size(), 2);
+    Assert.assertTrue(segments.contains("seg1"));
+    Assert.assertTrue(segments.contains("seg2"));
+  }
+
+  @Test
+  public void testGetSegmentToInstanceMap() {
+    Map<String, Map<String, String>> resultMap = 
idealStateHelper.getSegmentToInstanceMap("table");
+    Assert.assertEquals(resultMap.size(), 2);
+    Assert.assertEquals(resultMap.get("seg1").size(), 2);
+    Assert.assertEquals(resultMap.get("seg1").get("host1"), "ONLINE");
+    Assert.assertEquals(resultMap.get("seg1").get("host2"), "ONLINE");
+    Assert.assertEquals(resultMap.get("seg2").size(), 2);
+    Assert.assertEquals(resultMap.get("seg2").get("host2"), "ONLINE");
+    Assert.assertEquals(resultMap.get("seg2").get("host3"), "OFFLINE");
+  }
+}
\ No newline at end of file
diff --git a/pinot-grigio/pinot-grigio-coordinator/pom.xml 
b/pinot-grigio/pinot-grigio-coordinator/pom.xml
index f8f6e48..a7d6db2 100644
--- a/pinot-grigio/pinot-grigio-coordinator/pom.xml
+++ b/pinot-grigio/pinot-grigio-coordinator/pom.xml
@@ -89,6 +89,11 @@
       <version>${project.version}</version>
     </dependency>
     <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
       <groupId>org.testng</groupId>
       <artifactId>testng</artifactId>
       <scope>test</scope>
diff --git 
a/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogRetentionManagerImpl.java
 
b/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogRetentionManagerImpl.java
index ef13300..8c846b3 100644
--- 
a/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogRetentionManagerImpl.java
+++ 
b/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogRetentionManagerImpl.java
@@ -22,7 +22,7 @@ import org.apache.helix.NotificationContext;
 import org.apache.helix.api.listeners.IdealStateChangeListener;
 import org.apache.helix.model.IdealState;
 import org.apache.pinot.common.Utils;
-import org.apache.pinot.common.utils.IdealStateHelper;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
 import org.apache.pinot.grigio.common.storageProvider.UpdateLogStorageProvider;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
diff --git 
a/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogTableRetentionManagerImpl.java
 
b/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogTableRetentionManagerImpl.java
index cfb2f27..8e89fb9 100644
--- 
a/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogTableRetentionManagerImpl.java
+++ 
b/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogTableRetentionManagerImpl.java
@@ -18,40 +18,31 @@
  */
 package org.apache.pinot.grigio.common.storageProvider.retentionManager;
 
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
-import org.apache.pinot.common.utils.IdealStateHelper;
-import org.apache.pinot.common.utils.LLCSegmentName;
 import org.apache.pinot.grigio.common.storageProvider.UpdateLogStorageProvider;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 
-public class KCUpdateLogTableRetentionManagerImpl implements 
UpdateLogTableRetentionManager {
+public class KCUpdateLogTableRetentionManagerImpl extends 
UpdateLogTableRetentionManager {
 
   private static final Logger LOGGER = 
LoggerFactory.getLogger(KCUpdateLogTableRetentionManagerImpl.class);
 
   private Set<String> _segments;
-  private Map<Integer, LLCSegmentName> _partitionToLastSegment;
-  private final Map<String, String> _blackListedSegments;
-  private final String _tableName;
   private final UpdateLogStorageProvider _provider;
-  private final IdealStateHelper _idealStateHelper;
 
   public KCUpdateLogTableRetentionManagerImpl(IdealStateHelper 
idealStateHelper, String tableName, UpdateLogStorageProvider provider) throws 
IOException {
-
+    super(idealStateHelper, tableName);
     _provider = provider;
-    _tableName = tableName;
-    _idealStateHelper = idealStateHelper;
-    _blackListedSegments = new ConcurrentHashMap<>();
+    // load the current update log on this server and match it with helix 
stored state,
+    // so we can remove any unused update logs
     _provider.loadTable(tableName);
     _segments = _provider.getAllSegments(tableName);
-    _partitionToLastSegment = 
UpdateLogTableRetentionManager.getPartitionToLastSegment(_segments);
+    updateStateFromHelix();
   }
 
   public synchronized void updateSegmentsAndRemoveOldFiles(Set<String> 
newSegmentSet) {
@@ -69,40 +60,17 @@ public class KCUpdateLogTableRetentionManagerImpl 
implements UpdateLogTableReten
     }
   }
 
-  private synchronized void updateSegments(Set<String> newSegmentSet) {
-    _segments = newSegmentSet;
-    _partitionToLastSegment = 
UpdateLogTableRetentionManager.getPartitionToLastSegment(_segments);
-  }
-
-  public Set<String> getAllSegments() {
-    return ImmutableSet.copyOf(_segments);
-  }
-
   @Override
-  public synchronized boolean shouldIngestForSegment(String segmentName) {
-    if (_segments.contains(segmentName)) {
-      return true;
-    } else if (_blackListedSegments.containsKey(segmentName)) {
-      return false;
-    } else {
-      LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName);
-      int partition = llcSegmentName.getPartitionId();
-      if (!_partitionToLastSegment.containsKey(llcSegmentName.getPartitionId())
-          || UpdateLogTableRetentionManager.compareSegment(llcSegmentName, 
_partitionToLastSegment.get(partition))) {
-        updateStateFromHelix();
-        if (_segments.contains(segmentName)) {
-          LOGGER.info("segment {} matched in ideal state after refresh", 
segmentName);
-          return true;
-        }
-      }
-      _blackListedSegments.put(segmentName, segmentName);
-      return false;
-    }
+  protected boolean isSegmentAssignedToCurrentServer(String segmentName) {
+    // always return true as key coordinator should store whatever segment 
sent to it
+    return true;
   }
 
-  private void updateStateFromHelix() {
+  @Override
+  protected void updateStateFromHelix() {
     long start = System.currentTimeMillis();
-    
updateSegmentsAndRemoveOldFiles(_idealStateHelper.getSegmentToInstanceMap(_tableName).keySet());
+    super.updateStateFromHelix();
+    updateSegmentsAndRemoveOldFiles(_segmentsToInstanceMap.keySet());
     LOGGER.info("updated table {} state from helix in {} ms", _tableName, 
System.currentTimeMillis() - start);
   }
 
diff --git 
a/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/keyCoordinator/starter/KeyCoordinatorStarter.java
 
b/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/keyCoordinator/starter/KeyCoordinatorStarter.java
index 60012fc..c7d9628 100644
--- 
a/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/keyCoordinator/starter/KeyCoordinatorStarter.java
+++ 
b/pinot-grigio/pinot-grigio-coordinator/src/main/java/org/apache/pinot/grigio/keyCoordinator/starter/KeyCoordinatorStarter.java
@@ -24,7 +24,7 @@ import org.apache.commons.configuration.Configuration;
 import org.apache.commons.lang.StringUtils;
 import org.apache.pinot.common.metrics.MetricsHelper;
 import org.apache.pinot.common.utils.CommonConstants;
-import org.apache.pinot.common.utils.IdealStateHelper;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
 import org.apache.pinot.grigio.common.config.CommonConfig;
 import org.apache.pinot.grigio.common.rpcQueue.KeyCoordinatorQueueConsumer;
 import org.apache.pinot.grigio.common.rpcQueue.LogCoordinatorQueueProducer;
diff --git 
a/pinot-grigio/pinot-grigio-coordinator/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogRetentionManagerImplTest.java
 
b/pinot-grigio/pinot-grigio-coordinator/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogRetentionManagerImplTest.java
new file mode 100644
index 0000000..9f9f4de
--- /dev/null
+++ 
b/pinot-grigio/pinot-grigio-coordinator/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogRetentionManagerImplTest.java
@@ -0,0 +1,114 @@
+/**
+ * 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.pinot.grigio.common.storageProvider.retentionManager;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.apache.helix.NotificationContext;
+import org.apache.helix.model.IdealState;
+import org.apache.pinot.grigio.common.storageProvider.UpdateLogStorageProvider;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.mockito.internal.verification.VerificationModeFactory.times;
+
+public class KCUpdateLogRetentionManagerImplTest {
+  private KCUpdateLogRetentionManagerImpl kcUpdateLogRetentionManager;
+  private IdealStateHelper mockIdealStateHelper;
+  private UpdateLogStorageProvider mockUpdateLogStorageProvider;
+  private Map<String, Map<String, String>> table1SegmentMap;
+  private Map<String, Map<String, String>> table2SegmentMap;
+
+  @BeforeMethod
+  public void setUp() {
+    table1SegmentMap = new HashMap<>();
+
+    table1SegmentMap.put("table1__0__10__20191027T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server2", "ONLINE"));
+    table1SegmentMap.put("table1__0__11__20191028T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server3", "ONLINE"));
+    table1SegmentMap.put("table1__1__10__20191027T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server2", "ONLINE"));
+    table1SegmentMap.put("table1__2__10__20191027T2041Z",
+        ImmutableMap.of("server2", "ONLINE", "server3", "ONLINE"));
+
+    table2SegmentMap = new HashMap<>();
+    table2SegmentMap.put("table2__0__10__20191027T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server2", "ONLINE"));
+
+    mockIdealStateHelper = mock(IdealStateHelper.class);
+    
when(mockIdealStateHelper.getSegmentToInstanceMap("table1")).thenReturn(table1SegmentMap);
+    
when(mockIdealStateHelper.getSegmentToInstanceMap("table2")).thenReturn(table1SegmentMap);
+
+    mockUpdateLogStorageProvider = mock(UpdateLogStorageProvider.class);
+    kcUpdateLogRetentionManager = new 
KCUpdateLogRetentionManagerImpl(mockIdealStateHelper, 
mockUpdateLogStorageProvider,
+        "kc1");
+  }
+
+  @Test
+  public void testOnIdealStateChange() throws InterruptedException {
+    IdealState table1MockIdealState = getMockIdealState("table1_REALTIME",
+        ImmutableList.of("seg1", "host1"), ImmutableList.of("seg1", "host2"));
+
+    IdealState table2MockIdealState = getMockIdealState("table2_REALTIME",
+        ImmutableList.of("seg1", "host1"), ImmutableList.of("seg1", "host2"));
+
+    UpdateLogTableRetentionManager table1RetentionManager = 
kcUpdateLogRetentionManager.
+        getRetentionManagerForTable("table1_REALTIME");
+    UpdateLogTableRetentionManager table2RetentionManager = 
kcUpdateLogRetentionManager.
+        getRetentionManagerForTable("table2_REALTIME");
+
+    verify(mockIdealStateHelper, 
times(1)).getSegmentToInstanceMap("table1_REALTIME");
+    verify(mockIdealStateHelper, 
times(1)).getSegmentToInstanceMap("table2_REALTIME");
+
+    
kcUpdateLogRetentionManager.onIdealStateChange(ImmutableList.of(table1MockIdealState),
 mock(NotificationContext.class));
+    verify(mockIdealStateHelper, 
times(2)).getSegmentToInstanceMap("table1_REALTIME");
+    verify(mockIdealStateHelper, 
times(1)).getSegmentToInstanceMap("table2_REALTIME");
+
+    // table1 idealstate is the same as last time
+    
kcUpdateLogRetentionManager.onIdealStateChange(ImmutableList.of(table1MockIdealState,
 table2MockIdealState), mock(NotificationContext.class));
+    verify(mockIdealStateHelper, 
times(2)).getSegmentToInstanceMap("table1_REALTIME");
+    verify(mockIdealStateHelper, 
times(2)).getSegmentToInstanceMap("table2_REALTIME");
+
+    kcUpdateLogRetentionManager.onIdealStateChange(ImmutableList.of(
+        getMockIdealState("table1_REALTIME", ImmutableList.of("seg1", 
"host1"), ImmutableList.of("seg2", "host2")),
+        getMockIdealState("table2_REALTIME", ImmutableList.of("seg1", 
"host1"), ImmutableList.of("seg2", "host2"))
+        ), mock(NotificationContext.class));
+    verify(mockIdealStateHelper, 
times(3)).getSegmentToInstanceMap("table1_REALTIME");
+    verify(mockIdealStateHelper, 
times(3)).getSegmentToInstanceMap("table2_REALTIME");
+
+  }
+
+  private IdealState getMockIdealState(String resourceName, List<String>... 
partitionInstance) {
+    IdealState mockIdealState = new IdealState(resourceName);
+    mockIdealState.enable(true);
+    mockIdealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
+    Arrays.stream(partitionInstance).forEach(partition -> 
mockIdealState.setPartitionState(partition.get(0), partition.get(1), "ONLINE"));
+    return mockIdealState;
+  }
+}
\ No newline at end of file
diff --git 
a/pinot-grigio/pinot-grigio-coordinator/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogTableRetentionManagerImplTest.java
 
b/pinot-grigio/pinot-grigio-coordinator/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogTableRetentionManagerImplTest.java
new file mode 100644
index 0000000..c36ff1f
--- /dev/null
+++ 
b/pinot-grigio/pinot-grigio-coordinator/src/test/java/org/apache/pinot/grigio/common/storageProvider/retentionManager/KCUpdateLogTableRetentionManagerImplTest.java
@@ -0,0 +1,148 @@
+/**
+ * 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.pinot.grigio.common.storageProvider.retentionManager;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.pinot.grigio.common.storageProvider.UpdateLogStorageProvider;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class KCUpdateLogTableRetentionManagerImplTest {
+  private IdealStateHelper mockIdealStateHelper;
+  private UpdateLogStorageProvider mockStorageProvider;
+  private Map<String, Map<String, String>> segmentsInstanceMap;
+  private Set<String> updateLogStorageSegments;
+  private KCUpdateLogTableRetentionManagerImpl updateLogTableRetentionManager;
+
+  @BeforeMethod
+  public void setUp() throws IOException {
+    segmentsInstanceMap = new HashMap<>();
+
+    segmentsInstanceMap.put("table__0__10__20191027T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server2", "ONLINE"));
+    segmentsInstanceMap.put("table__0__11__20191028T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server3", "ONLINE"));
+    segmentsInstanceMap.put("table__1__10__20191027T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server2", "ONLINE"));
+    segmentsInstanceMap.put("table__2__10__20191027T2041Z",
+        ImmutableMap.of("server2", "ONLINE", "server3", "ONLINE"));
+
+    updateLogStorageSegments = new HashSet<>();
+    updateLogStorageSegments.add("table__0__10__20191027T2041Z");
+    updateLogStorageSegments.add("table__0__11__20191028T2041Z");
+    updateLogStorageSegments.add("table__1__10__20191027T2041Z");
+    updateLogStorageSegments.add("table__2__10__20191027T2041Z");
+
+    mockIdealStateHelper = mock(IdealStateHelper.class);
+    
when(mockIdealStateHelper.getSegmentToInstanceMap("table")).thenReturn(segmentsInstanceMap);
+
+    mockStorageProvider = mock(UpdateLogStorageProvider.class);
+    
when(mockStorageProvider.getAllSegments(anyString())).thenReturn(updateLogStorageSegments);
+
+  }
+
+  @Test
+  public void testInit() throws IOException {
+    updateLogStorageSegments.add("table__3__10__20191027T2041Z");
+    updateLogTableRetentionManager = new 
KCUpdateLogTableRetentionManagerImpl(mockIdealStateHelper, "table",
+        mockStorageProvider);
+
+    // verify that we update helix and performed deletion for expired segments
+    verify(mockStorageProvider, times(1))
+        .removeSegment("table", "table__3__10__20191027T2041Z");
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+  }
+
+  @Test
+  public void testShouldIngestForSegment() throws IOException {
+    updateLogTableRetentionManager = new 
KCUpdateLogTableRetentionManagerImpl(mockIdealStateHelper, "table",
+        mockStorageProvider);
+
+    // test the happy paths
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__0__10__20191027T2041Z"));
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__0__11__20191028T2041Z"));
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__1__10__20191027T2041Z"));
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__2__10__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    // test for segments that has lower seq
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__1__5__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    //test for segment that are older and not in existing segment list
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__1__20__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    //test for older timestamp but higher seq
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__0__20__20191027T2041Z"));
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+
+    // test for newer segments, and there is update in newer ideal state
+    segmentsInstanceMap.put("table__0__12__20191029T2041Z",
+        ImmutableMap.of("server1", "ONLINE", "server3", "ONLINE"));
+    
Assert.assertTrue(updateLogTableRetentionManager.shouldIngestForSegment("table__0__12__20191029T2041Z"));
+    verify(mockIdealStateHelper, times(2)).getSegmentToInstanceMap("table");
+
+    // test for newer segments and there is no update in newer ideal state
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__0__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(3)).getSegmentToInstanceMap("table");
+
+    // multiple attempts trying to fetch for segment should not recheck the 
ideal state often
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__0__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(3)).getSegmentToInstanceMap("table");
+
+    // for unknown partition, we should also refresh data
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__2__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(4)).getSegmentToInstanceMap("table");
+
+    
Assert.assertFalse(updateLogTableRetentionManager.shouldIngestForSegment("table__3__13__20191029T2141Z"));
+    verify(mockIdealStateHelper, times(5)).getSegmentToInstanceMap("table");
+
+  }
+
+  @Test
+  public void testNotifySegmentsChange() throws IOException {
+    updateLogTableRetentionManager = new 
KCUpdateLogTableRetentionManagerImpl(mockIdealStateHelper, "table",
+        mockStorageProvider);
+    verify(mockIdealStateHelper, times(1)).getSegmentToInstanceMap("table");
+    updateLogTableRetentionManager.notifySegmentsChange();
+    verify(mockIdealStateHelper, times(2)).getSegmentToInstanceMap("table");
+    verify(mockStorageProvider, never()).removeSegment(anyString(), 
anyString());
+
+    segmentsInstanceMap.remove("table__0__10__20191027T2041Z");
+    updateLogTableRetentionManager.notifySegmentsChange();
+    verify(mockIdealStateHelper, times(3)).getSegmentToInstanceMap("table");
+    verify(mockStorageProvider, times(1)).removeSegment("table", 
"table__0__10__20191027T2041Z");
+  }
+}
\ No newline at end of file
diff --git 
a/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
 
b/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
index c5c6886..d215cbd 100644
--- 
a/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
+++ 
b/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
@@ -24,7 +24,7 @@ import org.apache.helix.HelixManager;
 import org.apache.helix.ZNRecord;
 import org.apache.helix.store.zk.ZkHelixPropertyStore;
 import org.apache.pinot.common.metrics.ServerMetrics;
-import org.apache.pinot.common.utils.IdealStateHelper;
+import org.apache.pinot.grigio.common.utils.IdealStateHelper;
 import org.apache.pinot.core.data.manager.InstanceDataManager;
 import org.apache.pinot.core.query.executor.QueryExecutor;
 import org.apache.pinot.core.query.scheduler.QueryScheduler;


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to