diff --git a/gradle.properties b/gradle.properties index f3cc093a4c..01b578341c 100644 --- a/gradle.properties +++ b/gradle.properties @@ -14,7 +14,7 @@ kotlin.code.style=official # This is version for PgJdbc itself # Note: it should not include "-SNAPSHOT" as it is automatically added by build.gradle.kts # Release version can be generated by using -Prelease or -Prc= arguments -pgjdbc.version=42.7.3-yb-6 +pgjdbc.version=42.7.3-yb-7 lastEditYear=2024 # The options below configures the use of local clone (e.g. testing development versions) diff --git a/pgjdbc/build.gradle.kts b/pgjdbc/build.gradle.kts index f3da8c6875..5dc3b3fdad 100644 --- a/pgjdbc/build.gradle.kts +++ b/pgjdbc/build.gradle.kts @@ -91,6 +91,10 @@ dependencies { testFixturesImplementation("junit:junit:4.13.2") testFixturesImplementation("org.junit.jupiter:junit-jupiter-api:5.10.2") testFixturesImplementation("org.checkerframework:checker-qual:3.42.0") + + // 4.x is the last line with Java 8 bytecode; the CI matrix still runs tests on Java 8. + // mockito-junit-jupiter is not needed: the tests use Mockito.mock()/when() directly. + testImplementation("org.mockito:mockito-core:4.11.0") } val skipReplicationTests by props() diff --git a/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceService.java b/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceService.java index 89e38a7249..1e14199792 100644 --- a/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceService.java +++ b/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceService.java @@ -20,9 +20,15 @@ public class LoadBalanceService { static final byte STRICT_PREFERENCE = 0b00000001; protected static final String GET_SERVERS_QUERY = "select * from yb_servers()"; + // Control connection only runs yb_servers(). Keep these bounded so a hung refresh + // cannot stall synchronized checkAndRefresh() (and all new LB connections). + private static final int CONTROL_CONN_CONNECT_TIMEOUT_SECS = 10; + private static final int CONTROL_CONN_SOCKET_TIMEOUT_SECS = 15; + private static final int CONTROL_CONN_QUERY_TIMEOUT_SECS = 10; protected static final Logger LOGGER = Logger.getLogger("org.postgresql." + LoadBalanceService.class.getName()); private static boolean forceRefreshOnce = false; + private static volatile boolean loggedNoUniverseUuid = false; public static Map uuidToClusterInfoMap = new ConcurrentHashMap<>(); public static Map lbKeyToUuidMap = new ConcurrentHashMap<>(); @@ -60,6 +66,7 @@ static synchronized void clear() throws SQLException { } } forceRefreshOnce = false; + loggedNoUniverseUuid = false; } static long getLastRefreshTime(LoadBalancer lb) { @@ -88,17 +95,26 @@ private static boolean needsRefresh(long refreshInterval, LoadBalancer lb) { * @return the value of the universe_uuid of the cluster * @throws SQLException */ - private static synchronized String refresh(Connection conn, long refreshInterval, + static synchronized String refresh(Connection conn, long refreshInterval, LoadBalancer lb) throws SQLException { forceRefreshOnce = false; Statement st = conn.createStatement(); + st.setQueryTimeout(CONTROL_CONN_QUERY_TIMEOUT_SECS); LOGGER.fine("Executing query: " + GET_SERVERS_QUERY + " to fetch list of servers"); ResultSet rs = st.executeQuery(GET_SERVERS_QUERY); InetAddress hostConnectedInetAddress = getConnectedInetAddress(conn); ConcurrentHashMap hostToNodeInfoMap = null; boolean publicIPsGivenForAll = true; + // A public_ip we cannot even resolve is not a usable connection target. Tracked separately + // from publicIPsGivenForAll so the two reasons to reject public addresses stay distinct. + boolean publicIPsResolvableForAll = true; String uuid = null; ClusterInfo cluster = null; + // Case-sensitive, to match hostToNodeInfoMap. Comparing case-insensitively here while the + // map keys case-sensitively lets a stale entry survive eviction forever: a host returned + // with different case is inserted under the new spelling, while the old spelling looks + // "still present" and is never removed. + Set hostsInLatestQuery = new HashSet<>(); while (rs.next()) { String host = rs.getString("host"); LOGGER.finest("Received entry for host " + host); @@ -112,9 +128,12 @@ private static synchronized String refresh(Connection conn, long refreshInterval try { uuid = rs.getString("universe_uuid"); } catch (PSQLException e) { - LOGGER.info("Found a version of YugabyteDB which does not send universe_uuid in its " - + "response of yb_servers() function. Connecting to more than one such clusters is not " - + "supported."); + if (!loggedNoUniverseUuid) { + LOGGER.info("Found a version of YugabyteDB which does not send universe_uuid in its " + + "response of yb_servers() function. Connecting to more than one such clusters is not " + + "supported."); + loggedNoUniverseUuid = true; + } uuid = "default"; } @@ -130,14 +149,18 @@ private static synchronized String refresh(Connection conn, long refreshInterval hostToNodeInfoMap = cluster.getHostToNodeInfoMap() != null ? cluster.getHostToNodeInfoMap() : new ConcurrentHashMap<>(); } - NodeInfo nodeInfo = hostToNodeInfoMap.containsKey(host) ? hostToNodeInfoMap.get(host) : - new NodeInfo(); + // Key this node the same way the map is currently keyed, so that the lookup below, the + // insert further down and the eviction after the loop all address the same entry. A + // previous refresh may have re-keyed the map by public_ip. + String key = cluster.isKeyedByPublicIp() && !publicHost.isEmpty() ? publicHost : host; + hostsInLatestQuery.add(key); + NodeInfo nodeInfo = hostToNodeInfoMap.containsKey(key) ? hostToNodeInfoMap.get(key) + : new NodeInfo(); synchronized (nodeInfo) { nodeInfo.host = host; nodeInfo.publicIP = publicHost; - publicIPsGivenForAll = !publicHost.isEmpty(); + publicIPsGivenForAll = publicIPsGivenForAll ? !publicHost.isEmpty() : false; nodeInfo.placement = new CloudPlacement(cloud, region, zone); - LOGGER.fine("Setting node_type to " + nodeType + " for host " + host); nodeInfo.nodeType = nodeType; try { nodeInfo.port = Integer.valueOf(port); @@ -148,30 +171,37 @@ private static synchronized String refresh(Connection conn, long refreshInterval long failedHostTTL = Long.getLong(FAILED_HOST_RECONNECT_DELAY_SECS_KEY, DEFAULT_FAILED_HOST_TTL_SECONDS); if (nodeInfo.isDown) { if (System.currentTimeMillis() - nodeInfo.isDownSince > (failedHostTTL * 1000)) { - LOGGER.fine("Marking " + nodeInfo.host + " as UP since failed-host-reconnect-delay-secs (" + failedHostTTL + "s) has elapsed"); + LOGGER.info("Marked " + nodeInfo.host + " as UP"); nodeInfo.isDown = false; } else { - LOGGER.fine("Keeping " + nodeInfo.host + " as DOWN since failed-host-reconnect-delay-secs (" + failedHostTTL + "s) has not elapsed"); + LOGGER.fine("Kept " + nodeInfo.host + " marked as DOWN since failed-host-reconnect-delay-secs (" + failedHostTTL + "s) has not elapsed"); } } } - hostToNodeInfoMap.putIfAbsent(host, nodeInfo); + if (hostToNodeInfoMap.putIfAbsent(key, nodeInfo) == null) { + // Only when the node is new to the map. Logging every node on every refresh would + // swamp the log, since a failed connection forces a refresh via forceRefreshOnce. + LOGGER.info("Added " + nodeInfo + " to the host map"); + } InetAddress hostInetAddr; InetAddress publicHostInetAddr; try { hostInetAddr = InetAddress.getByName(host); } catch (UnknownHostException e) { - LOGGER.fine("Failed to get host '" + host + "' by name"); + LOGGER.fine("Failed to get inet address for host '" + host + "' by name"); hostInetAddr = null; } try { publicHostInetAddr = !publicHost.isEmpty() ? InetAddress.getByName(publicHost) : null; } catch (UnknownHostException e) { - LOGGER.fine("Failed to get public_ip '" + publicHost + "' by name"); + LOGGER.fine("Failed to get inet address for public_ip '" + publicHost + "' by name"); publicHostInetAddr = null; } + if (!publicHost.isEmpty() && publicHostInetAddr == null) { + publicIPsResolvableForAll = false; + } if (cluster.getUseHostColumn() == null) { if (hostConnectedInetAddress.equals(hostInetAddr)) { cluster.setUseHostColumn(Boolean.TRUE); @@ -181,6 +211,27 @@ private static synchronized String refresh(Connection conn, long refreshInterval // Both host and public_ip are same cluster.setUseHostColumn(Boolean.TRUE); } + if (cluster.getUseHostColumn() != null) { + LOGGER.info("Will use '" + (cluster.getUseHostColumn() ? "host" : "public_ip") + + "' addresses for connections"); + } + } + } + + if (cluster == null || uuid == null || hostToNodeInfoMap == null) { + LOGGER.warning("yb_servers() returned no rows, skipping refresh"); + // Still advance the refresh clock. Leaving it alone keeps needsRefresh() true, so every + // subsequent getConnection() would run another yb_servers() under the class monitor. + lb.setLastRefreshTime(System.currentTimeMillis()); + return lb.getUuid(); + } + + Set removedHosts = new HashSet<>(hostToNodeInfoMap.keySet()); + removedHosts.removeAll(hostsInLatestQuery); + if (!removedHosts.isEmpty()) { + LOGGER.info("Evicting hosts no longer returned by yb_servers(): " + removedHosts); + for (String h : removedHosts) { + hostToNodeInfoMap.remove(h); } } @@ -190,17 +241,39 @@ private static synchronized String refresh(Connection conn, long refreshInterval } lb.setUuid(uuid); } - if ((cluster.getUseHostColumn() != null && !cluster.getUseHostColumn()) || (cluster.getUseHostColumn() == null && publicIPsGivenForAll)) { - LOGGER.info("Will be using publicIPs for establishing connections"); - ArrayList hosts = Collections.list(hostToNodeInfoMap.keys()); - for (String host : hosts) { - NodeInfo info = hostToNodeInfoMap.get(host); - hostToNodeInfoMap.remove(info.host); - hostToNodeInfoMap.put(info.publicIP, info); + // With useHostColumn undecided, public addresses are a guess: only take it when every node + // gave one and every one of them resolves. Otherwise the balancer would hand out addresses + // it cannot connect to and fall back to unbalanced connections for the life of the pool. + // useHostColumn == FALSE is not a guess -- it means the control connection reached the + // cluster on a public address -- so it is honoured as before. + boolean publicIpsUnusable = cluster.getUseHostColumn() == null && publicIPsGivenForAll + && !publicIPsResolvableForAll; + boolean usePublicIp = (cluster.getUseHostColumn() != null && !cluster.getUseHostColumn()) + || (cluster.getUseHostColumn() == null && publicIPsGivenForAll + && publicIPsResolvableForAll); + if (usePublicIp) { + cluster.setLoggedUnusablePublicIps(false); + if (!cluster.isKeyedByPublicIp()) { + LOGGER.info("Re-keying the host map by 'public_ip' addresses"); + rekeyBy(hostToNodeInfoMap, true); + cluster.setKeyedByPublicIp(true); + } + } else { + if (cluster.isKeyedByPublicIp()) { + LOGGER.info("Re-keying the host map by 'host' addresses"); + rekeyBy(hostToNodeInfoMap, false); + cluster.setKeyedByPublicIp(false); + } + if (publicIpsUnusable) { + if (!cluster.hasLoggedUnusablePublicIps()) { + LOGGER.warning("Not using 'public_ip' addresses: they are set for all nodes but could " + + "not be resolved. Using 'host' addresses instead."); + cluster.setLoggedUnusablePublicIps(true); + } + } else if (cluster.getUseHostColumn() == null) { + LOGGER.warning("Unable to identify set of addresses to use for establishing connections. " + + "Using 'host' addresses."); } - } else if (cluster.getUseHostColumn() == null) { - LOGGER.warning("Unable to identify set of addresses to use for establishing connections. " - + "Using private addresses."); } lb.setLastRefreshTime(System.currentTimeMillis()); if (cluster != null) { @@ -210,21 +283,70 @@ private static synchronized String refresh(Connection conn, long refreshInterval return uuid; } + /** + * Re-keys every entry by its public_ip (or back by its host). Removes by the key actually + * iterated rather than by a field of the value, so an entry cannot be dropped twice and a + * stale key cannot survive. A node without a public_ip stays keyed by its host, which is the + * same fallback the per-row key uses. + */ + private static void rekeyBy(ConcurrentHashMap hostToNodeInfoMap, + boolean publicIp) { + for (String key : Collections.list(hostToNodeInfoMap.keys())) { + NodeInfo info = hostToNodeInfoMap.get(key); + if (info == null) { + continue; + } + String newKey = publicIp && !info.publicIP.isEmpty() ? info.publicIP : info.host; + if (!newKey.equals(key)) { + hostToNodeInfoMap.remove(key); + hostToNodeInfoMap.put(newKey, info); + } + } + } + + /** + * The configured timeout for this metadata-only path, capped at capSecs. A value of 0 disables + * the timeout in pgjdbc, and absent means the driver default, so both take the cap. + */ + private static String cappedTimeoutSecs(Properties properties, String name, int capSecs) { + String configured = properties.getProperty(name); + if (configured != null && !configured.trim().isEmpty()) { + try { + int secs = Integer.parseInt(configured.trim()); + if (secs > 0 && secs < capSecs) { + return String.valueOf(secs); + } + } catch (NumberFormatException nfe) { + LOGGER.fine("Could not parse " + name + " '" + configured + "', using " + capSecs + "s" + + " for the control connection"); + } + } + return String.valueOf(capSecs); + } + private static void markAsFailed(String uuid, String host) { if (uuid == null) { - LOGGER.fine("Could not mark down host: " + host + ". UUID not present"); + LOGGER.info("Could not mark host: " + host + " as DOWN. Cluster UUID not present"); return; } NodeInfo info = uuidToClusterInfoMap.get(uuid).getHostToNodeInfoMap().get(host); if (info == null) { - return; // unexpected + // Expected for the cluster endpoint: the control connection may be dialing the address + // from the URL, which is not a node and so never a key in the host map. + LOGGER.fine("Could not mark host: " + host + " as DOWN. Host not found in cluster UUID: " + + uuid); + return; } synchronized (info) { String previous = info.isDown ? "DOWN" : "UP"; info.isDown = true; info.isDownSince = System.currentTimeMillis(); info.connectionCount = 0; - LOGGER.info("Marked " + host + " as DOWN (was " + previous + " earlier)"); + if (previous.equals("UP")) { + LOGGER.info("Marked " + host + " as DOWN (was UP earlier)"); + } else { + LOGGER.fine(host + " is already marked as DOWN"); + } } } @@ -359,7 +481,7 @@ private static LoadBalancer getLB(LoadBalanceProperties.LoadBalancerKey key) { // 2. create control connection and fetch yb_servers() -- refer to LBProperties.checkAndRefresh() String uuid = checkAndRefresh(key, lb); if (uuid == null) { - LOGGER.fine("Attempt to refresh info from yb_servers() failed"); + LOGGER.warning("Attempt to refresh info from yb_servers() failed"); return null; } @@ -402,14 +524,14 @@ private static Connection getConnection(LoadBalanceProperties.LoadBalancerKey ke } LoadBalancer lb = getLB(key); if (lb == null) { - LOGGER.fine("No loadbalancer found for lbkey: " + key); + LOGGER.warning("No loadbalancer found for lbkey: " + key); return null; } String uuid = lbKeyToUuidMap.get(key); String url = key.getUrl(); if (checkAndRefresh(key, lb) == null) { - LOGGER.fine("Attempt to refresh info from yb_servers() failed"); + LOGGER.warning("Attempt to refresh info from yb_servers() failed"); return null; } @@ -471,7 +593,13 @@ private static synchronized String checkAndRefresh(LoadBalanceProperties.LoadBal if (needsRefresh(lb.getRefreshListSeconds(), lb)) { String url = key.getUrl(); Properties properties = new Properties(key.getProperties()); - properties.setProperty("socketTimeout", "15"); + // Cap, do not override: a disabled (0), absent or oversized app timeout must not be + // inherited by this metadata-only path, but a user asking for something shorter than the + // cap gets what they asked for. These are awaited under the class monitor. + properties.setProperty("connectTimeout", + cappedTimeoutSecs(properties, "connectTimeout", CONTROL_CONN_CONNECT_TIMEOUT_SECS)); + properties.setProperty("socketTimeout", + cappedTimeoutSecs(properties, "socketTimeout", CONTROL_CONN_SOCKET_TIMEOUT_SECS)); // Preserve the host from the URL for TLS hostname verification. The first attempt below // dials this host, but the retry path points the control connection at a node address // from yb_servers(), which a cluster-wide certificate does not cover. @@ -483,16 +611,20 @@ private static synchronized String checkAndRefresh(LoadBalanceProperties.LoadBal + LoadBalanceProperties.ENDPOINT_HOST_KEY + " for TLS hostname verification of the control connection"); } - HostSpec[] hspec = hostSpecs(properties); + // The configured host(s) from the URL. hspec is pointed at fetched node addresses as the + // loop advances, so keep the original to be able to come back to it. + HostSpec[] configuredHspec = hostSpecs(properties); + HostSpec[] hspec = configuredHspec; Connection controlConnection = null; ArrayList hosts = getAllAvailableHosts(lbKeyToUuidMap.get(key), new ArrayList<>()); + boolean configuredHostsRetried = false; while (true) { boolean refreshFailed = false; try { if (uuid != null) { controlConnection = uuidToClusterInfoMap.get(uuid).getControlConnection(); } - if (controlConnection == null){ + if (controlConnection == null) { controlConnection = new PgConnection(hspec, properties, url); if (uuid != null) { uuidToClusterInfoMap.get(uuid).setControlConnection(controlConnection); @@ -508,14 +640,18 @@ private static synchronized String checkAndRefresh(LoadBalanceProperties.LoadBal } } catch (SQLException ex) { if (refreshFailed) { - LOGGER.fine("Exception while refreshing: " + ex + ", " + ex.getSQLState()); + LOGGER.warning("Exception while refreshing: " + ex + ", " + ex.getSQLState()); String failed = ((PgConnection) controlConnection).getQueryExecutor().getHostSpec().getHost(); - markAsFailed(uuid, failed); + // A yb_servers() that ran past CONTROL_CONN_QUERY_TIMEOUT_SECS says the cluster is + // slow, not that this node refuses client connections, so do not mark it DOWN. + if (!PSQLState.QUERY_CANCELED.getState().equals(ex.getSQLState())) { + markAsFailed(uuid, failed); + } // Drop the host we just failed to refresh against so we don't retry it indefinitely. hosts.remove(failed); } else { String msg = hspec.length > 1 ? " and others" : ""; - LOGGER.fine("Exception while creating control connection to " + LOGGER.warning("Exception while creating control connection to " + hspec[0].getHost() + msg + ": " + ex + ", " + ex.getSQLState()); for (HostSpec h : hspec) { hosts.remove(h.getHost()); @@ -540,17 +676,26 @@ private static synchronized String checkAndRefresh(LoadBalanceProperties.LoadBal " YugabyteDB, consider upgrading it."); return null; } - // Retry until servers are available - if (hosts.isEmpty()) { - LOGGER.warning("Failed to establish control connection to available servers"); - return null; - } else { + // Where to dial next. On a refresh failure, come back to the configured host(s) once + // before moving on to fetched node addresses. Worth doing even when hspec already + // points at them: the failed connection was usually the cached one rather than one + // dialed in this pass, and a configured host is often a service or load balancer + // address, so re-dialing it lands on whichever pod is currently healthy. + if (refreshFailed && !configuredHostsRetried) { + LOGGER.warning("Refresh failed. Retrying with the configured host(s): " + + configuredHspec[0].getHost() + (configuredHspec.length > 1 ? " and others" : "")); + hspec = configuredHspec; + configuredHostsRetried = true; + } else if (!hosts.isEmpty()) { // Try the first host in the list (don't have to check least loaded one since it's // just for the control connection). This also advances off a host whose refresh // failed, so we don't keep hammering the same node. HostSpec hs = new HostSpec(hosts.get(0), getPort(uuid, hosts.get(0)), key.getProperties().getProperty("localSocketAddress")); hspec = new HostSpec[]{hs}; + } else { + LOGGER.warning("Failed to establish control connection to available servers"); + return null; } } } @@ -649,6 +794,11 @@ public String getNodeType() { public void setNodeType(String nodeType) { this.nodeType = nodeType; } + + public String toString() { + return "NodeInfo=[" + host + "," + port + "," + publicIP + "," + nodeType + "," + placement + + "," + connectionCount + "," + isDown + "," + isDownSince + "]"; + } } static class CloudPlacement { @@ -713,6 +863,17 @@ public static class ClusterInfo { private Map lbKeyToLBMap = new ConcurrentHashMap<>(); private Boolean useHostColumn = null; + /** + * Which key form {@link #hostToNodeInfoMap} is currently in. The tail of + * {@link #refresh} re-keys the map by public_ip for clusters reached that way, so lookups, + * inserts and eviction in the next refresh must all use the same form. + */ + private volatile boolean keyedByPublicIp = false; + /** + * Latches the "public_ip set but unresolvable" warning so it is logged on entering that + * state rather than on every refresh. Per cluster, so one cluster cannot mute another. + */ + private volatile boolean loggedUnusablePublicIps = false; public Connection getControlConnection() { return controlConnection; @@ -734,6 +895,22 @@ public Map getLbKeyToLBMap( return lbKeyToLBMap; } + public boolean hasLoggedUnusablePublicIps() { + return loggedUnusablePublicIps; + } + + public void setLoggedUnusablePublicIps(boolean loggedUnusablePublicIps) { + this.loggedUnusablePublicIps = loggedUnusablePublicIps; + } + + public boolean isKeyedByPublicIp() { + return keyedByPublicIp; + } + + public void setKeyedByPublicIp(boolean keyedByPublicIp) { + this.keyedByPublicIp = keyedByPublicIp; + } + public Boolean getUseHostColumn() { return useHostColumn; } diff --git a/pgjdbc/src/test/java/com/yugabyte/ysql/LoadBalanceServiceRefreshTest.java b/pgjdbc/src/test/java/com/yugabyte/ysql/LoadBalanceServiceRefreshTest.java new file mode 100644 index 0000000000..5f08775c4c --- /dev/null +++ b/pgjdbc/src/test/java/com/yugabyte/ysql/LoadBalanceServiceRefreshTest.java @@ -0,0 +1,505 @@ +package com.yugabyte.ysql; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import org.postgresql.jdbc.PgConnection; +import org.postgresql.util.PSQLException; +import org.postgresql.util.PSQLState; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; + +class LoadBalanceServiceRefreshTest { + + private PgConnection mockConn; + private Statement mockStmt; + private ResultSet mockRs; + private TestLoadBalancer lb; + + @BeforeEach + void setUp() throws SQLException { + LoadBalanceService.uuidToClusterInfoMap.clear(); + LoadBalanceService.lbKeyToUuidMap.clear(); + + mockConn = mock(PgConnection.class, RETURNS_DEEP_STUBS); + when(mockConn.getQueryExecutor().getHostSpec().getHost()).thenReturn("127.0.0.1"); + + mockStmt = mock(Statement.class); + mockRs = mock(ResultSet.class); + when(mockConn.createStatement()).thenReturn(mockStmt); + when(mockStmt.executeQuery(anyString())).thenReturn(mockRs); + + lb = new TestLoadBalancer(); + } + + @AfterEach + void tearDown() throws SQLException { + LoadBalanceService.clear(); + } + + @Test + void emptyResultSetReturnsEarlyWithoutNPE() throws SQLException { + when(mockRs.next()).thenReturn(false); + lb.setUuid("existing-uuid"); + + String result = assertDoesNotThrow( + () -> LoadBalanceService.refresh(mockConn, 300, lb)); + assertEquals("existing-uuid", result); + assertTrue(LoadBalanceService.uuidToClusterInfoMap.isEmpty()); + } + + @Test + void emptyResultSetWithNullUuidReturnsNull() throws SQLException { + when(mockRs.next()).thenReturn(false); + + String result = assertDoesNotThrow( + () -> LoadBalanceService.refresh(mockConn, 300, lb)); + assertNull(result); + } + + @Test + void removedHostsAreDetected() throws SQLException { + String uuid = "test-uuid"; + ConcurrentHashMap hostMap = new ConcurrentHashMap<>(); + addNodeInfo(hostMap, "127.0.0.1", "10.0.0.1", 5433, "aws", "us-west", "us-west-2a"); + addNodeInfo(hostMap, "127.0.0.2", "10.0.0.2", 5433, "aws", "us-west", "us-west-2b"); + addNodeInfo(hostMap, "127.0.0.3", "10.0.0.3", 5433, "aws", "us-west", "us-west-2c"); + + LoadBalanceService.ClusterInfo cluster = new LoadBalanceService.ClusterInfo(); + cluster.setHostToNodeInfoMap(hostMap); + cluster.setUseHostColumn(Boolean.TRUE); + cluster.setControlConnection(mockConn); + LoadBalanceService.uuidToClusterInfoMap.put(uuid, cluster); + + setupResultSetRows( + row("127.0.0.1", "10.0.0.1", "5433", "aws", "us-west", "us-west-2a", "primary", uuid), + row("127.0.0.2", "10.0.0.2", "5433", "aws", "us-west", "us-west-2b", "primary", uuid) + ); + + lb.setUuid(uuid); + LoadBalanceService.refresh(mockConn, 300, lb); + + ConcurrentHashMap updatedMap = + LoadBalanceService.uuidToClusterInfoMap.get(uuid).getHostToNodeInfoMap(); + assertTrue(updatedMap.containsKey("127.0.0.1")); + assertTrue(updatedMap.containsKey("127.0.0.2")); + // 127.0.0.3 is absent because it's not returned by yb_servers() function + assertFalse(updatedMap.containsKey("127.0.0.3")); + } + + @Test + void longDownRemovedHostsAreEvicted() throws SQLException { + String uuid = "test-uuid"; + ConcurrentHashMap hostMap = new ConcurrentHashMap<>(); + addNodeInfo(hostMap, "127.0.0.1", "10.0.0.1", 5433, "aws", "us-west", "us-west-2a"); + addNodeInfo(hostMap, "127.0.0.2", "10.0.0.2", 5433, "aws", "us-west", "us-west-2b"); + + LoadBalanceService.NodeInfo downNode = + addNodeInfo(hostMap, "127.0.0.3", "10.0.0.3", 5433, "aws", "us-west", "us-west-2c"); + markAsDown(downNode, System.currentTimeMillis() - 400_000); + + LoadBalanceService.ClusterInfo cluster = new LoadBalanceService.ClusterInfo(); + cluster.setHostToNodeInfoMap(hostMap); + cluster.setUseHostColumn(Boolean.TRUE); + cluster.setControlConnection(mockConn); + LoadBalanceService.uuidToClusterInfoMap.put(uuid, cluster); + + setupResultSetRows( + row("127.0.0.1", "10.0.0.1", "5433", "aws", "us-west", "us-west-2a", "primary", uuid), + row("127.0.0.2", "10.0.0.2", "5433", "aws", "us-west", "us-west-2b", "primary", uuid) + ); + + lb.setUuid(uuid); + LoadBalanceService.refresh(mockConn, 300, lb); + + ConcurrentHashMap updatedMap = + LoadBalanceService.uuidToClusterInfoMap.get(uuid).getHostToNodeInfoMap(); + assertTrue(updatedMap.containsKey("127.0.0.1")); + assertTrue(updatedMap.containsKey("127.0.0.2")); + assertFalse(updatedMap.containsKey("127.0.0.3"), + "Host down for >300s and absent from yb_servers() should be evicted"); + } + + // SKIP, because we remove stale host immediately + void recentlyDownRemovedHostsAreNotEvicted() throws SQLException { + String uuid = "test-uuid"; + ConcurrentHashMap hostMap = new ConcurrentHashMap<>(); + addNodeInfo(hostMap, "127.0.0.1", "10.0.0.1", 5433, "aws", "us-west", "us-west-2a"); + + LoadBalanceService.NodeInfo downNode = + addNodeInfo(hostMap, "127.0.0.2", "10.0.0.2", 5433, "aws", "us-west", "us-west-2b"); + markAsDown(downNode, System.currentTimeMillis() - 60_000); + + LoadBalanceService.ClusterInfo cluster = new LoadBalanceService.ClusterInfo(); + cluster.setHostToNodeInfoMap(hostMap); + cluster.setUseHostColumn(Boolean.TRUE); + cluster.setControlConnection(mockConn); + LoadBalanceService.uuidToClusterInfoMap.put(uuid, cluster); + + setupResultSetRows( + row("127.0.0.1", "10.0.0.1", "5433", "aws", "us-west", "us-west-2a", "primary", uuid) + ); + + lb.setUuid(uuid); + LoadBalanceService.refresh(mockConn, 300, lb); + + ConcurrentHashMap updatedMap = + LoadBalanceService.uuidToClusterInfoMap.get(uuid).getHostToNodeInfoMap(); + assertTrue(updatedMap.containsKey("127.0.0.2"), + "Host down for <300s should NOT be evicted even if absent from yb_servers()"); + } + + /** + * hostToNodeInfoMap is case-sensitive, so a host returned under a different spelling is a + * different key. It replaces the old entry rather than being matched to it: comparing + * case-insensitively while keying case-sensitively would insert the new spelling and leave + * the old one looking "still present", so it could never be evicted. + */ + @Test + void hostReturnedWithDifferentCaseReplacesStaleEntry() throws SQLException { + String uuid = "test-uuid"; + ConcurrentHashMap hostMap = new ConcurrentHashMap<>(); + addNodeInfo(hostMap, "HostA", "10.0.0.1", 5433, "aws", "us-west", "us-west-2a"); + + LoadBalanceService.NodeInfo downNode = + addNodeInfo(hostMap, "HostB", "10.0.0.2", 5433, "aws", "us-west", "us-west-2b"); + markAsDown(downNode, System.currentTimeMillis() - 400_000); + + LoadBalanceService.ClusterInfo cluster = new LoadBalanceService.ClusterInfo(); + cluster.setHostToNodeInfoMap(hostMap); + cluster.setUseHostColumn(Boolean.TRUE); + cluster.setControlConnection(mockConn); + LoadBalanceService.uuidToClusterInfoMap.put(uuid, cluster); + + // yb_servers() returns "hosta" (lowercase); "HostB" is gone from the cluster entirely. + setupResultSetRows( + row("hosta", "10.0.0.1", "5433", "aws", "us-west", "us-west-2a", "primary", uuid) + ); + + lb.setUuid(uuid); + LoadBalanceService.refresh(mockConn, 300, lb); + + ConcurrentHashMap updatedMap = + LoadBalanceService.uuidToClusterInfoMap.get(uuid).getHostToNodeInfoMap(); + assertTrue(updatedMap.containsKey("hosta"), "the spelling yb_servers() returned is kept"); + assertFalse(updatedMap.containsKey("HostA"), + "the stale spelling must be evicted, not left behind as a duplicate of the same node"); + assertFalse(updatedMap.containsKey("HostB"), "a host no longer returned must be evicted"); + assertEquals(1, updatedMap.size(), "one entry per node returned by yb_servers()"); + } + + /** + * Public-IP cluster: yb_servers() reports the private address in "host" and the routable + * address in "public_ip", so the previous refresh left the map keyed by public_ip. A refresh + * that still reports the node must not reset its connection count. + */ + @Test + void publicIpKeyedMapRetainsConnectionCountAcrossRefresh() throws SQLException { + String uuid = "test-uuid"; + seedPublicIpKeyedCluster(uuid); + + LoadBalanceService.refresh(mockConn, 300, lb); + + ConcurrentHashMap updatedMap = + LoadBalanceService.uuidToClusterInfoMap.get(uuid).getHostToNodeInfoMap(); + assertEquals(2, updatedMap.size(), "map should hold one entry per node, keyed by public_ip"); + assertTrue(updatedMap.containsKey("node-a.invalid")); + assertEquals(5, updatedMap.get("node-a.invalid").getConnectionCount(), + "connection count must survive a refresh that still reports the node"); + } + + /** + * Same setup, for the DOWN state: a node marked down stays down until + * failed-host-reconnect-delay-secs elapses, otherwise the balancer immediately re-picks a node + * it just failed to reach. + */ + @Test + void publicIpKeyedMapRetainsDownStateAcrossRefresh() throws SQLException { + String uuid = "test-uuid"; + seedPublicIpKeyedCluster(uuid); + + LoadBalanceService.refresh(mockConn, 300, lb); + + ConcurrentHashMap updatedMap = + LoadBalanceService.uuidToClusterInfoMap.get(uuid).getHostToNodeInfoMap(); + assertTrue(updatedMap.containsKey("node-b.invalid")); + assertTrue(updatedMap.get("node-b.invalid").isDown(), + "a host marked DOWN must stay DOWN until failed-host-reconnect-delay-secs elapses"); + } + + /** + * The driver could not decide which address set to use: the control connection is to an + * endpoint that is neither the node's host nor its public_ip (a k8s service or load balancer). + * Every node gives a public_ip and every one resolves, so public addresses are used and node + * state survives the refresh. + */ + @Test + void undeterminedHostColumnWithResolvablePublicIpsKeepsPublicIpKeys() throws SQLException { + String uuid = "test-uuid"; + seedPublicIpKeyedCluster(uuid, "10.1.0.1", "10.1.0.2"); + LoadBalanceService.uuidToClusterInfoMap.get(uuid).setUseHostColumn(null); + + LoadBalanceService.refresh(mockConn, 300, lb); + + LoadBalanceService.ClusterInfo cluster = LoadBalanceService.uuidToClusterInfoMap.get(uuid); + ConcurrentHashMap updatedMap = + cluster.getHostToNodeInfoMap(); + assertTrue(cluster.isKeyedByPublicIp()); + assertTrue(updatedMap.containsKey("10.1.0.1"), "map should stay keyed by public_ip"); + assertEquals(5, updatedMap.get("10.1.0.1").getConnectionCount(), + "connection count must survive a refresh that still reports the node"); + } + + /** + * Same, except the public_ip values are set but do not resolve (k8s-internal names seen from + * outside the cluster). Guessing public addresses there hands the balancer targets it cannot + * connect to, so the map falls back to host addresses -- carrying node state with it. + */ + @Test + void undeterminedHostColumnWithUnresolvablePublicIpsFallsBackToHost() throws SQLException { + String uuid = "test-uuid"; + seedPublicIpKeyedCluster(uuid); + LoadBalanceService.uuidToClusterInfoMap.get(uuid).setUseHostColumn(null); + + LoadBalanceService.refresh(mockConn, 300, lb); + + LoadBalanceService.ClusterInfo cluster = LoadBalanceService.uuidToClusterInfoMap.get(uuid); + ConcurrentHashMap updatedMap = + cluster.getHostToNodeInfoMap(); + assertFalse(cluster.isKeyedByPublicIp()); + assertTrue(updatedMap.containsKey("10.0.0.1"), "map should fall back to host addresses"); + assertFalse(updatedMap.containsKey("node-a.invalid"), + "unresolvable public_ip must not be left as a key"); + assertEquals(2, updatedMap.size(), "one entry per node"); + assertEquals(5, updatedMap.get("10.0.0.1").getConnectionCount(), + "connection count must survive the fallback re-key"); + } + + @Test + void missingUniverseUuidFallsBackToDefault() throws SQLException { + when(mockRs.next()).thenReturn(true, false); + when(mockRs.getString("host")).thenReturn("127.0.0.1"); + when(mockRs.getString("public_ip")).thenReturn(""); + when(mockRs.getString("port")).thenReturn("5433"); + when(mockRs.getString("cloud")).thenReturn("aws"); + when(mockRs.getString("region")).thenReturn("us-west"); + when(mockRs.getString("zone")).thenReturn("us-west-2a"); + when(mockRs.getString("node_type")).thenReturn("primary"); + when(mockRs.getString("universe_uuid")).thenThrow( + new PSQLException("column not found", PSQLState.UNDEFINED_COLUMN)); + + String result = LoadBalanceService.refresh(mockConn, 300, lb); + assertEquals("default", result); + assertTrue(LoadBalanceService.uuidToClusterInfoMap.containsKey("default")); + } + + // ---- helpers ---- + + private void setupResultSetRows(String[]... rows) throws SQLException { + if (rows.length == 0) { + when(mockRs.next()).thenReturn(false); + return; + } + Boolean[] nextReturns = new Boolean[rows.length + 1]; + for (int i = 0; i < rows.length; i++) { + nextReturns[i] = true; + } + nextReturns[rows.length] = false; + when(mockRs.next()).thenReturn(nextReturns[0], + java.util.Arrays.copyOfRange(nextReturns, 1, nextReturns.length)); + + String[] hosts = new String[rows.length]; + String[] publicIps = new String[rows.length]; + String[] ports = new String[rows.length]; + String[] clouds = new String[rows.length]; + String[] regions = new String[rows.length]; + String[] zones = new String[rows.length]; + String[] nodeTypes = new String[rows.length]; + String[] uuids = new String[rows.length]; + for (int i = 0; i < rows.length; i++) { + hosts[i] = rows[i][0]; + publicIps[i] = rows[i][1]; + ports[i] = rows[i][2]; + clouds[i] = rows[i][3]; + regions[i] = rows[i][4]; + zones[i] = rows[i][5]; + nodeTypes[i] = rows[i][6]; + uuids[i] = rows[i][7]; + } + setupGetString("host", hosts); + setupGetString("public_ip", publicIps); + setupGetString("port", ports); + setupGetString("cloud", clouds); + setupGetString("region", regions); + setupGetString("zone", zones); + setupGetString("node_type", nodeTypes); + setupGetString("universe_uuid", uuids); + } + + private void setupGetString(String column, String[] values) throws SQLException { + if (values.length == 1) { + when(mockRs.getString(column)).thenReturn(values[0]); + } else { + when(mockRs.getString(column)).thenReturn(values[0], + java.util.Arrays.copyOfRange(values, 1, values.length)); + } + } + + private static String[] row(String host, String publicIp, String port, String cloud, + String region, String zone, String nodeType, String uuid) { + return new String[]{host, publicIp, port, cloud, region, zone, nodeType, uuid}; + } + + private static LoadBalanceService.NodeInfo addNodeInfo( + ConcurrentHashMap map, + String host, String publicIp, int port, String cloud, String region, String zone) { + LoadBalanceService.NodeInfo info = new LoadBalanceService.NodeInfo(); + try { + java.lang.reflect.Field hostField = LoadBalanceService.NodeInfo.class.getDeclaredField("host"); + hostField.setAccessible(true); + hostField.set(info, host); + java.lang.reflect.Field publicIPField = LoadBalanceService.NodeInfo.class.getDeclaredField("publicIP"); + publicIPField.setAccessible(true); + publicIPField.set(info, publicIp); + java.lang.reflect.Field portField = LoadBalanceService.NodeInfo.class.getDeclaredField("port"); + portField.setAccessible(true); + portField.setInt(info, port); + java.lang.reflect.Field placementField = LoadBalanceService.NodeInfo.class.getDeclaredField("placement"); + placementField.setAccessible(true); + placementField.set(info, new LoadBalanceService.CloudPlacement(cloud, region, zone)); + java.lang.reflect.Field nodeTypeField = LoadBalanceService.NodeInfo.class.getDeclaredField("nodeType"); + nodeTypeField.setAccessible(true); + nodeTypeField.set(info, "primary"); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + map.put(host, info); + return info; + } + + private static void markAsDown(LoadBalanceService.NodeInfo info, long downSince) { + try { + java.lang.reflect.Field isDownField = LoadBalanceService.NodeInfo.class.getDeclaredField("isDown"); + isDownField.setAccessible(true); + isDownField.setBoolean(info, true); + java.lang.reflect.Field isDownSinceField = LoadBalanceService.NodeInfo.class.getDeclaredField("isDownSince"); + isDownSinceField.setAccessible(true); + isDownSinceField.setLong(info, downSince); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + } + + /** + * Seeds a two-node cluster in the state a previous refresh leaves behind on a public-IP + * cluster: map keyed by public_ip, useHostColumn FALSE. Node A carries 5 connections, node B + * was just marked DOWN. yb_servers() then reports both nodes as still present. + */ + private void seedPublicIpKeyedCluster(String uuid) throws SQLException { + seedPublicIpKeyedCluster(uuid, "node-a.invalid", "node-b.invalid"); + } + + private void seedPublicIpKeyedCluster(String uuid, String publicIpA, String publicIpB) + throws SQLException { + ConcurrentHashMap hostMap = new ConcurrentHashMap<>(); + LoadBalanceService.NodeInfo nodeA = + addNodeInfo(hostMap, "10.0.0.1", publicIpA, 5433, "aws", "us-west", "us-west-2a"); + LoadBalanceService.NodeInfo nodeB = + addNodeInfo(hostMap, "10.0.0.2", publicIpB, 5433, "aws", "us-west", "us-west-2b"); + keyByPublicIp(hostMap, nodeA); + keyByPublicIp(hostMap, nodeB); + + setConnectionCount(nodeA, 5); + markAsDown(nodeB, System.currentTimeMillis()); + + LoadBalanceService.ClusterInfo cluster = new LoadBalanceService.ClusterInfo(); + cluster.setHostToNodeInfoMap(hostMap); + cluster.setUseHostColumn(Boolean.FALSE); + // The previous refresh re-keyed the map by public_ip and recorded that on the cluster. + cluster.setKeyedByPublicIp(true); + cluster.setControlConnection(mockConn); + LoadBalanceService.uuidToClusterInfoMap.put(uuid, cluster); + + setupResultSetRows( + row("10.0.0.1", publicIpA, "5433", "aws", "us-west", "us-west-2a", "primary", uuid), + row("10.0.0.2", publicIpB, "5433", "aws", "us-west", "us-west-2b", "primary", uuid) + ); + lb.setUuid(uuid); + } + + /** Re-keys an entry by its public_ip, as the tail of refresh() does for public-IP clusters. */ + private static void keyByPublicIp(ConcurrentHashMap map, + LoadBalanceService.NodeInfo info) { + map.remove(info.getHost()); + map.put(info.getPublicIP(), info); + } + + private static void setConnectionCount(LoadBalanceService.NodeInfo info, int count) { + try { + java.lang.reflect.Field field = + LoadBalanceService.NodeInfo.class.getDeclaredField("connectionCount"); + field.setAccessible(true); + field.setInt(info, count); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + } + + private static class TestLoadBalancer implements LoadBalancer { + private String uuid; + private long lastRefreshTime; + + @Override + public boolean isHostEligible(java.util.Map.Entry e, + Byte requestFlags) { + return true; + } + + @Override + public String getLeastLoadedServer(boolean newRequest, List failedHosts, + ArrayList timedOutHosts) { + return null; + } + + @Override + public int getRefreshListSeconds() { + return 300; + } + + @Override + public void setUuid(String uuid) { + this.uuid = uuid; + } + + @Override + public String getUuid() { + return uuid; + } + + @Override + public long getLastRefreshTime() { + return lastRefreshTime; + } + + @Override + public void setLastRefreshTime(long lastRefreshTime) { + this.lastRefreshTime = lastRefreshTime; + } + } +}