From 102fc92120bbc9d22a7aac53a99645270963a407 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Tue, 1 Mar 2016 17:20:33 -0800
Subject: [PATCH] SQLMetadataConnector: Fix overzealous retries that were
 preventing EntryExistsException from making it out.

---
 .../druid/metadata/SQLMetadataConnector.java  | 24 ++++++++++++++-----
 .../SQLMetadataStorageActionHandler.java      | 12 ++++++++++
 .../SQLMetadataStorageActionHandlerTest.java  | 17 +++++++++++++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/io/druid/metadata/SQLMetadataConnector.java b/server/src/main/java/io/druid/metadata/SQLMetadataConnector.java
index a004e69927c..160a9ba320a 100644
--- a/server/src/main/java/io/druid/metadata/SQLMetadataConnector.java
+++ b/server/src/main/java/io/druid/metadata/SQLMetadataConnector.java
@@ -100,7 +100,10 @@ public abstract class SQLMetadataConnector implements MetadataStorageConnector
 
   public abstract boolean tableExists(Handle handle, final String tableName);
 
-  public <T> T retryWithHandle(final HandleCallback<T> callback)
+  public <T> T retryWithHandle(
+      final HandleCallback<T> callback,
+      final Predicate<Throwable> myShouldRetry
+  )
   {
     final Callable<T> call = new Callable<T>()
     {
@@ -112,13 +115,18 @@ public abstract class SQLMetadataConnector implements MetadataStorageConnector
     };
     final int maxTries = 10;
     try {
-      return RetryUtils.retry(call, shouldRetry, maxTries);
+      return RetryUtils.retry(call, myShouldRetry, maxTries);
     }
     catch (Exception e) {
       throw Throwables.propagate(e);
     }
   }
 
+  public <T> T retryWithHandle(final HandleCallback<T> callback)
+  {
+    return retryWithHandle(callback, shouldRetry);
+  }
+
   public <T> T retryTransaction(final TransactionCallback<T> callback)
   {
     final Callable<T> call = new Callable<T>()
@@ -399,21 +407,24 @@ public abstract class SQLMetadataConnector implements MetadataStorageConnector
   }
 
   @Override
-  public void createRulesTable() {
+  public void createRulesTable()
+  {
     if (config.get().isCreateTables()) {
       createRulesTable(tablesConfigSupplier.get().getRulesTable());
     }
   }
 
   @Override
-  public void createConfigTable() {
+  public void createConfigTable()
+  {
     if (config.get().isCreateTables()) {
       createConfigTable(tablesConfigSupplier.get().getConfigTable());
     }
   }
 
   @Override
-  public void createTaskTables() {
+  public void createTaskTables()
+  {
     if (config.get().isCreateTables()) {
       final MetadataStorageTablesConfig tablesConfig = tablesConfigSupplier.get();
       final String entryType = tablesConfig.getTaskEntryType();
@@ -502,7 +513,8 @@ public abstract class SQLMetadataConnector implements MetadataStorageConnector
     );
   }
   @Override
-  public void createAuditTable() {
+  public void createAuditTable()
+  {
     if (config.get().isCreateTables()) {
       createAuditTable(tablesConfigSupplier.get().getAuditTable());
     }
diff --git a/server/src/main/java/io/druid/metadata/SQLMetadataStorageActionHandler.java b/server/src/main/java/io/druid/metadata/SQLMetadataStorageActionHandler.java
index 9996157a398..1681a3f69b2 100644
--- a/server/src/main/java/io/druid/metadata/SQLMetadataStorageActionHandler.java
+++ b/server/src/main/java/io/druid/metadata/SQLMetadataStorageActionHandler.java
@@ -22,6 +22,7 @@ package io.druid.metadata;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Optional;
+import com.google.common.base.Predicate;
 import com.google.common.base.Throwables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
@@ -115,6 +116,17 @@ public class SQLMetadataStorageActionHandler<EntryType, StatusType, LogType, Loc
                     .execute();
               return null;
             }
+          },
+          new Predicate<Throwable>()
+          {
+            @Override
+            public boolean apply(Throwable e)
+            {
+              final boolean isStatementException = e instanceof StatementException ||
+                                                   (e instanceof CallbackFailedException
+                                                    && e.getCause() instanceof StatementException);
+              return connector.isTransientException(e) && !(isStatementException && getEntry(id).isPresent());
+            }
           }
       );
     }
diff --git a/server/src/test/java/io/druid/metadata/SQLMetadataStorageActionHandlerTest.java b/server/src/test/java/io/druid/metadata/SQLMetadataStorageActionHandlerTest.java
index 5ff26a2cc96..0547e47e49d 100644
--- a/server/src/test/java/io/druid/metadata/SQLMetadataStorageActionHandlerTest.java
+++ b/server/src/test/java/io/druid/metadata/SQLMetadataStorageActionHandlerTest.java
@@ -32,6 +32,7 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 
 import java.util.HashSet;
 import java.util.Map;
@@ -41,6 +42,9 @@ public class SQLMetadataStorageActionHandlerTest
   @Rule
   public final TestDerbyConnector.DerbyConnectorRule derbyConnectorRule = new TestDerbyConnector.DerbyConnectorRule();
 
+  @Rule
+  public final ExpectedException thrown = ExpectedException.none();
+
   private static final ObjectMapper jsonMapper = new DefaultObjectMapper();
   private SQLMetadataStorageActionHandler<Map<String, Integer>, Map<String, Integer>, Map<String, String>, Map<String, Integer>> handler;
 
@@ -176,6 +180,19 @@ public class SQLMetadataStorageActionHandlerTest
     );
   }
 
+  @Test(timeout = 10_000L)
+  public void testRepeatInsert() throws Exception
+  {
+    final String entryId = "abcd";
+    Map<String, Integer> entry = ImmutableMap.of("a", 1);
+    Map<String, Integer> status = ImmutableMap.of("count", 42);
+
+    handler.insert(entryId, new DateTime("2014-01-01"), "test", entry, true, status);
+
+    thrown.expect(EntryExistsException.class);
+    handler.insert(entryId, new DateTime("2014-01-01"), "test", entry, true, status);
+  }
+
   @Test
   public void testLogs() throws Exception
   {