Repository: prestodb/presto
Issue: Raptor local shard cleanup should only mark shards for current node

The calls in `ShardCleaner` to `updateCleanedShardNodes()` and `updatePurgedShardNodes()` should take the current node ID. It is updating the shard for all nodes rather than just the local node.


Changed files:
presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardManagerDao.java
presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardCleaner.java
presto-raptor/src/test/java/com/facebook/presto/raptor/metadata/TestShardManagerDao.java

diff --git a/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardCleaner.java b/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardCleaner.java
index b7ef3d1d4c..7138fa225e 100644
--- a/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardCleaner.java
+++ b/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardCleaner.java
@@ -16,6 +16,7 @@ package com.facebook.presto.raptor.metadata;
 import com.facebook.presto.raptor.backup.BackupStore;
 import com.facebook.presto.raptor.storage.StorageService;
 import com.facebook.presto.spi.NodeManager;
+import com.facebook.presto.spi.PrestoException;
 import com.google.common.annotations.VisibleForTesting;
 import io.airlift.log.Logger;
 import io.airlift.units.Duration;
@@ -41,6 +42,7 @@ import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 
+import static com.facebook.presto.raptor.RaptorErrorCode.RAPTOR_ERROR;
 import static com.facebook.presto.raptor.util.DatabaseUtil.onDemandDao;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
 import static java.util.Objects.requireNonNull;
@@ -239,7 +241,7 @@ public class ShardCleaner
             for (UUID uuid : uuids) {
                 deleteFile(storageService.getStorageFile(uuid));
             }
-            dao.updateCleanedShardNodes(uuids);
+            dao.updateCleanedShardNodes(uuids, getCurrentNodeId());
         }
     }
 
@@ -255,7 +257,7 @@ public class ShardCleaner
             for (UUID uuid : uuids) {
                 deleteFile(storageService.getStorageFile(uuid));
             }
-            dao.updatePurgedShardNodes(uuids);
+            dao.updatePurgedShardNodes(uuids, getCurrentNodeId());
         }
     }
 
@@ -327,6 +329,15 @@ public class ShardCleaner
         }
     }
 
+    private int getCurrentNodeId()
+    {
+        Integer nodeId = dao.getNodeId(currentNode);
+        if (nodeId == null) {
+            throw new PrestoException(RAPTOR_ERROR, "Node does not exist: " + currentNode);
+        }
+        return nodeId;
+    }
+
     private static Timestamp maxTimestamp(Duration duration)
     {
         return new Timestamp(System.currentTimeMillis() - duration.toMillis());
diff --git a/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardManagerDao.java b/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardManagerDao.java
index 98e3ee0b0c..e6ebdbf69c 100644
--- a/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardManagerDao.java
+++ b/presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardManagerDao.java
@@ -279,11 +279,17 @@ public interface ShardManagerDao
 
     @SqlBatch("UPDATE deleted_shard_nodes SET clean_time = CURRENT_TIMESTAMP\\n" +
             "WHERE shard_uuid = :shardUuid\\n" +
+            "  AND node_id = :nodeId\\n" +
             "  AND clean_time IS NULL\\n")
-    void updateCleanedShardNodes(@Bind("shardUuid") Iterable<UUID> shardUuids);
+    void updateCleanedShardNodes(
+            @Bind("shardUuid") Iterable<UUID> shardUuids,
+            @Bind("nodeId") int nodeId);
 
     @SqlBatch("UPDATE deleted_shard_nodes SET purge_time = CURRENT_TIMESTAMP\\n" +
             "WHERE shard_uuid = :shardUuid\\n" +
+            "  AND node_id = :nodeId\\n" +
             "  AND purge_time IS NULL\\n")
-    void updatePurgedShardNodes(@Bind("shardUuid") Iterable<UUID> shardUuids);
+    void updatePurgedShardNodes(
+            @Bind("shardUuid") Iterable<UUID> shardUuids,
+            @Bind("nodeId") int nodeId);
 }
diff --git a/presto-raptor/src/test/java/com/facebook/presto/raptor/metadata/TestShardManagerDao.java b/presto-raptor/src/test/java/com/facebook/presto/raptor/metadata/TestShardManagerDao.java
index eb56cdc0ef..2335e5ad09 100644
--- a/presto-raptor/src/test/java/com/facebook/presto/raptor/metadata/TestShardManagerDao.java
+++ b/presto-raptor/src/test/java/com/facebook/presto/raptor/metadata/TestShardManagerDao.java
@@ -24,12 +24,14 @@ import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
 import java.sql.SQLException;
+import java.sql.Timestamp;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 
 import static com.facebook.presto.raptor.metadata.SchemaDaoUtil.createTablesWithRetry;
 import static io.airlift.testing.Assertions.assertInstanceOf;
+import static java.util.concurrent.TimeUnit.DAYS;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
@@ -276,6 +278,50 @@ public class TestShardManagerDao
         assertContainsShardNode(shardNodes, nodeName2, shardUuid4);
     }
 
+    @Test
+    public void testDeletedShardNodes()
+    {
+        String nodeName1 = UUID.randomUUID().toString();
+        int nodeId1 = dao.insertNode(nodeName1);
+
+        String nodeName2 = UUID.randomUUID().toString();
+        int nodeId2 = dao.insertNode(nodeName2);
+
+        ImmutableList<UUID> shards = ImmutableList.of(UUID.randomUUID());
+
+        // insert shard on both nodes
+        dao.insertDeletedShardNodes(shards, ImmutableList.of(nodeId1));
+        dao.insertDeletedShardNodes(shards, ImmutableList.of(nodeId2));
+
+        // verify we should clean from both
+        assertEquals(dao.getCleanableShardNodesBatch(nodeName1, future()), shards);
+        assertEquals(dao.getCleanableShardNodesBatch(nodeName2, future()), shards);
+
+        // clean on first node
+        dao.updateCleanedShardNodes(shards, nodeId1);
+        assertEquals(dao.getCleanableShardNodesBatch(nodeName1, future()), ImmutableList.of());
+        assertEquals(dao.getCleanableShardNodesBatch(nodeName2, future()), shards);
+
+        // clean on second node
+        dao.updateCleanedShardNodes(shards, nodeId2);
+        assertEquals(dao.getCleanableShardNodesBatch(nodeName1, future()), ImmutableList.of());
+        assertEquals(dao.getCleanableShardNodesBatch(nodeName2, future()), ImmutableList.of());
+
+        // verify we should purge from both
+        assertEquals(dao.getPurgableShardNodesBatch(nodeName1, future()), shards);
+        assertEquals(dao.getPurgableShardNodesBatch(nodeName2, future()), shards);
+
+        // purge on first node
+        dao.updatePurgedShardNodes(shards, nodeId1);
+        assertEquals(dao.getPurgableShardNodesBatch(nodeName1, future()), ImmutableList.of());
+        assertEquals(dao.getPurgableShardNodesBatch(nodeName2, future()), shards);
+
+        // purge on second node
+        dao.updatePurgedShardNodes(shards, nodeId2);
+        assertEquals(dao.getPurgableShardNodesBatch(nodeName1, future()), ImmutableList.of());
+        assertEquals(dao.getPurgableShardNodesBatch(nodeName2, future()), ImmutableList.of());
+    }
+
     private long createTable(String name)
     {
         return dbi.onDemand(MetadataDao.class).insertTable("test", name, false);
@@ -285,4 +331,9 @@ public class TestShardManagerDao
     {
         assertTrue(nodes.contains(new ShardNode(shardUuid, nodeName)));
     }
+
+    private static Timestamp future()
+    {
+        return new Timestamp(System.currentTimeMillis() + DAYS.toMillis(1));
+    }
 }