Apply backoff when metastore connection fails

When createMetastoreClient fails to establish a connection, the failure
was only recorded as the last exception without updating the backoff
state, unlike API-level failures which already back off. As a result the
unreachable metastore was never penalized in the candidate ordering and
kept being tried first on every call, producing repeated connection
attempts to it instead of preferring a healthy fallback. Invoke
backoff.fail() on connection failure as well.

Co-authored-by: Marius Grama <findinpath@gmail.com>

Author: 2 people

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java

@@ -103,6 +103,7 @@ public ThriftMetastoreClient createMetastoreClient(Optional<String> delegationTo
             }
             catch (TException e) {
                 lastException = e;
+                backoff.fail();
             }
         }
 

plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestStaticTokenAwareMetastoreClientFactory.java

@@ -19,9 +19,12 @@
 import io.airlift.testing.TestingTicker;
 import io.trino.hive.thrift.metastore.Table;
 import org.apache.thrift.TException;
+import org.apache.thrift.transport.TTransportException;
 import org.junit.jupiter.api.Test;
 
 import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 
@@ -183,6 +186,27 @@ public void testReturnsToDefaultClientAfterErrorOnFallback()
         assertEqualHiveClient(metastoreClient3, DEFAULT_CLIENT);
     }
 
+    @Test
+    public void testBackoffAppliedOnConnectionFailure()
+            throws TException
+    {
+        ConnectionCountingThriftMetastoreClientFactory clientFactory = new ConnectionCountingThriftMetastoreClientFactory(
+                ImmutableMap.of(DEFAULT_URI, Optional.empty(), FALLBACK_URI, Optional.of(FALLBACK_CLIENT)));
+        StaticMetastoreConfig config = new StaticMetastoreConfig().setMetastoreUris(ImmutableList.of(DEFAULT_URI, FALLBACK_URI));
+        StaticTokenAwareMetastoreClientFactory metastoreClientFactory = new StaticTokenAwareMetastoreClientFactory(
+                config, new ThriftMetastoreAuthenticationConfig(), clientFactory, new TestingTicker());
+
+        assertEqualHiveClient(metastoreClientFactory.createMetastoreClient(Optional.empty()), FALLBACK_CLIENT);
+        assertThat(clientFactory.connectionAttempts()).containsExactlyInAnyOrderEntriesOf(ImmutableMap.of(
+                URI.create(DEFAULT_URI), 1,
+                URI.create(FALLBACK_URI), 1));
+        clientFactory.resetConnectionAttempts();
+        assertEqualHiveClient(metastoreClientFactory.createMetastoreClient(Optional.empty()), FALLBACK_CLIENT);
+        // The previous connection failure on DEFAULT_URI is not attempted first again on the second call
+        assertThat(clientFactory.connectionAttempts()).containsExactlyInAnyOrderEntriesOf(ImmutableMap.of(
+                URI.create(FALLBACK_URI), 1));
+    }
+
     private static void assertGetTableException(ThriftMetastoreClient client)
     {
         assertThatThrownBy(() -> client.getTable("foo", "bar"))
@@ -230,4 +254,34 @@ private static void assertEqualHiveClient(ThriftMetastoreClient actual, ThriftMe
         }
         assertThat(actual).isEqualTo(expected);
     }
+
+    private static class ConnectionCountingThriftMetastoreClientFactory
+            implements ThriftMetastoreClientFactory
+    {
+        private final ThriftMetastoreClientFactory delegate;
+        private final Map<URI, Integer> attempts = new HashMap<>();
+
+        public ConnectionCountingThriftMetastoreClientFactory(Map<String, Optional<ThriftMetastoreClient>> clients)
+        {
+            this.delegate = new MockThriftMetastoreClientFactory(clients);
+        }
+
+        @Override
+        public ThriftMetastoreClient create(URI uri, Optional<String> delegationToken)
+                throws TTransportException
+        {
+            attempts.merge(uri, 1, Integer::sum);
+            return delegate.create(uri, delegationToken);
+        }
+
+        public void resetConnectionAttempts()
+        {
+            attempts.clear();
+        }
+
+        public Map<URI, Integer> connectionAttempts()
+        {
+            return ImmutableMap.copyOf(attempts);
+        }
+    }
 }