Multibinding for RetryStrategy

trinodb · Dec 15, 2023 · 7cc22ac · 7cc22ac
1 parent 3aea15e
commit 7cc22ac
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 14 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactory.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactory.java
@@ -25,28 +25,30 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.sql.SQLTransientException;
+import java.util.Set;
 
 import static java.time.temporal.ChronoUnit.MILLIS;
 import static java.time.temporal.ChronoUnit.SECONDS;
+import static java.util.Collections.singleton;
 import static java.util.Objects.requireNonNull;
 
 public class RetryingConnectionFactory
         implements ConnectionFactory
 {
+    private final Set<RetryStrategy> retryStrategies;
     private final RetryPolicy<Object> retryPolicy;
-
     private final ConnectionFactory delegate;
 
     @Inject
-    public RetryingConnectionFactory(StatisticsAwareConnectionFactory delegate, RetryStrategy retryStrategy)
+    public RetryingConnectionFactory(StatisticsAwareConnectionFactory delegate, Set<RetryStrategy> retryStrategies)
     {
-        requireNonNull(retryStrategy);
+        this.retryStrategies = retryStrategiesWithDefault(requireNonNull(retryStrategies));
         this.delegate = requireNonNull(delegate, "delegate is null");
         this.retryPolicy = RetryPolicy.builder()
                 .withMaxDuration(java.time.Duration.of(30, SECONDS))
                 .withMaxAttempts(5)
                 .withBackoff(50, 5_000, MILLIS, 4)
-                .handleIf(retryStrategy::isExceptionRecoverable)
+                .handleIf(this::combineRetryStrategies)
                 .abortOn(TrinoException.class)
                 .build();
     }
@@ -89,4 +91,23 @@ public boolean isExceptionRecoverable(Throwable exception)
                     .anyMatch(SQLTransientException.class::isInstance);
         }
     }
+
+    private boolean combineRetryStrategies(Throwable throwable)
+    {
+        return retryStrategies.stream()
+                .map(retryStrategy -> retryStrategy.isExceptionRecoverable(throwable))
+                .filter(recoverable -> recoverable)
+                .findAny()
+                .orElse(false);
+    }
+
+    private static Set<RetryStrategy> retryStrategiesWithDefault(Set<RetryStrategy> retryStrategies)
+    {
+        if (retryStrategies.isEmpty()) {
+            return singleton(new DefaultRetryStrategy());
+        }
+        else {
+            return retryStrategies;
+        }
+    }
 }
diff --git a/...n/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactoryModule.java b/...n/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactoryModule.java
@@ -15,10 +15,9 @@
 
 import com.google.inject.AbstractModule;
 import com.google.inject.Scopes;
-import io.trino.plugin.jdbc.RetryingConnectionFactory.DefaultRetryStrategy;
 import io.trino.plugin.jdbc.RetryingConnectionFactory.RetryStrategy;
 
-import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+import static com.google.inject.multibindings.Multibinder.newSetBinder;
 
 public class RetryingConnectionFactoryModule
         extends AbstractModule
@@ -27,9 +26,6 @@ public class RetryingConnectionFactoryModule
     public void configure()
     {
         bind(RetryingConnectionFactory.class).in(Scopes.SINGLETON);
-        newOptionalBinder(binder(), RetryStrategy.class)
-                .setDefault()
-                .to(DefaultRetryStrategy.class)
-                .in(Scopes.SINGLETON);
+        newSetBinder(binder(), RetryStrategy.class);
     }
 }
diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestRetryingConnectionFactory.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestRetryingConnectionFactory.java
@@ -34,7 +34,7 @@
 import java.util.stream.Stream;
 
 import static com.google.common.reflect.Reflection.newProxy;
-import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.RETURN;
 import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_NPE;
 import static io.trino.plugin.jdbc.TestRetryingConnectionFactory.MockConnectorFactory.Action.THROW_SQL_EXCEPTION;
@@ -171,7 +171,7 @@ private static Injector createInjectorWithOverridenStrategy(MockConnectorFactory
             binder.bind(MockConnectorFactory.class).in(Scopes.SINGLETON);
             binder.bind(ConnectionFactory.class).annotatedWith(ForBaseJdbc.class).to(Key.get(MockConnectorFactory.class));
             binder.install(new RetryingConnectionFactoryModule());
-            newOptionalBinder(binder, RetryStrategy.class).setBinding().to(OverrideRetryStrategy.class).in(Scopes.SINGLETON);
+            newSetBinder(binder, RetryStrategy.class).addBinding().to(OverrideRetryStrategy.class);
         });
     }
 

diff --git a/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClientModule.java b/plugin/trino-oracle/src/main/java/io/trino/plugin/oracle/OracleClientModule.java
@@ -55,7 +55,7 @@ public void configure(Binder binder)
         configBinder(binder).bindConfig(OracleConfig.class);
         newOptionalBinder(binder, Key.get(int.class, MaxDomainCompactionThreshold.class)).setBinding().toInstance(ORACLE_MAX_LIST_EXPRESSIONS);
         newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(Query.class).in(Scopes.SINGLETON);
-        newOptionalBinder(binder, RetryStrategy.class).setBinding().to(OracleRetryStrategy.class).in(Scopes.SINGLETON);
+        newSetBinder(binder, RetryStrategy.class).addBinding().to(OracleRetryStrategy.class);
     }
 
     @Provides

diff --git a/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestingOracleServer.java b/plugin/trino-oracle/src/test/java/io/trino/plugin/oracle/TestingOracleServer.java
@@ -44,6 +44,7 @@
 import static io.trino.testing.TestingConnectorSession.SESSION;
 import static io.trino.testing.containers.TestContainers.startOrReuse;
 import static java.lang.String.format;
+import static java.util.Collections.singleton;
 
 public class TestingOracleServer
         implements AutoCloseable
@@ -131,7 +132,7 @@ private ConnectionFactory getConnectionFactory(String connectionUrl, String user
                 new OracleDriver(),
                 new BaseJdbcConfig().setConnectionUrl(connectionUrl),
                 StaticCredentialProvider.of(username, password)));
-        return new RetryingConnectionFactory(connectionFactory, new DefaultRetryStrategy());
+        return new RetryingConnectionFactory(connectionFactory, singleton(new DefaultRetryStrategy()));
     }
 
     @Override