[#9788] YSQL SMARTDRIVER: Switch to using smart driver for yb-pgsql u…

…nit tests

Summary:
Currently the yb-pgsql test uses the upstream PostgreSQL driver version 42.2.5. Changing it to use the latest 'beta' version of the jdbc-yugabytedb version which is 42.3.0-beta.1. The mvn dependency is:

```
<dependency>
      <groupId>com.yugabyte</groupId>
      <artifactId>jdbc-yugabytedb</artifactId>
      <version>42.3.0-beta.1</version>
    </dependency>
```

Please note the upstream version is pgjdbc 42.3.0.

Additionally the unit tests for load balance has been augmented to test whether the load is close to evenly distributed when multiple connections are being created concurrently.
This test also uses the exact port of servers for connectivity rather than hard code value of '5433' used earlier. This is corresponding to the change done in commit revision c4308a8

Test Plan:
The testLoadBalance unit test is augmented to take care of the above mentioned changes.
The entire unit test was run locally against the new version and they ran fine.

Reviewers: zyu, mihnea

Reviewed By: zyu, mihnea

Subscribers: yql

Differential Revision: https://phabricator.dev.yugabyte.com/D12689

java/yb-pgsql/pom.xml

@@ -28,8 +28,9 @@
       <scope>test</scope>
     </dependency>
     <dependency>
-      <groupId>org.postgresql</groupId>
-      <artifactId>postgresql</artifactId>
+      <groupId>com.yugabyte</groupId>
+      <artifactId>jdbc-yugabytedb</artifactId>
+      <version>42.3.0</version>
     </dependency>
     <dependency>
       <groupId>org.json</groupId>

java/yb-pgsql/src/test/java/org/yb/pgsql/BasePgSQLTest.java

@@ -29,11 +29,11 @@
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
-import org.postgresql.core.TransactionState;
-import org.postgresql.jdbc.PgArray;
-import org.postgresql.jdbc.PgConnection;
-import org.postgresql.util.PGobject;
-import org.postgresql.util.PSQLException;
+import com.yugabyte.core.TransactionState;
+import com.yugabyte.jdbc.PgArray;
+import com.yugabyte.jdbc.PgConnection;
+import com.yugabyte.util.PGobject;
+import com.yugabyte.util.PSQLException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.yb.client.IsInitDbDoneResponse;
@@ -266,7 +266,7 @@ public void initPostgresBefore() throws Exception {
       return;
 
     LOG.info("Loading PostgreSQL JDBC driver");
-    Class.forName("org.postgresql.Driver");
+    Class.forName("com.yugabyte.Driver");
 
     // Postgres bin directory.
     pgBinDir = new File(TestUtils.getBuildRootDir(), "postgres/bin");
@@ -312,7 +312,7 @@ protected void resetSettings() {
     startRedisProxy = false;
   }
 
-  static ConnectionBuilder getConnectionBuilder() {
+  protected ConnectionBuilder getConnectionBuilder() {
     return new ConnectionBuilder(miniCluster);
   }
 
@@ -1769,6 +1769,7 @@ public static class ConnectionBuilder implements Cloneable {
     private static final int INITIAL_CONNECTION_DELAY_MS = 500;
 
     private final MiniYBCluster miniCluster;
+    private boolean loadBalance;
 
     private int tserverIndex = 0;
     private String database = DEFAULT_PG_DATABASE;
@@ -1865,7 +1866,7 @@ Connection connect() throws Exception {
       final InetSocketAddress postgresAddress = miniCluster.getPostgresContactPoints()
           .get(tserverIndex);
       String url = String.format(
-          "jdbc:postgresql://%s:%d/%s",
+          "jdbc:yugabytedb://%s:%d/%s",
           postgresAddress.getHostName(),
           postgresAddress.getPort(),
           database
@@ -1895,6 +1896,9 @@ Connection connect() throws Exception {
         props.setProperty("loggerLevel", "TRACE");
       }
 
+      boolean loadBalance = getLoadBalance();
+      String lbValue = loadBalance ? "true" : "false";
+      props.setProperty("load-balance", lbValue);
       int delayMs = INITIAL_CONNECTION_DELAY_MS;
       for (int attempt = 1; attempt <= MAX_CONNECTION_ATTEMPTS; ++attempt) {
         Connection connection = null;
@@ -1950,5 +1954,13 @@ Connection connect() throws Exception {
       }
       throw new IllegalStateException("Should not be able to reach here");
     }
+
+    public boolean getLoadBalance() {
+      return loadBalance;
+    }
+
+    public void setLoadBalance(boolean lb) {
+      loadBalance = lb;
+    }
   }
 }

java/yb-pgsql/src/test/java/org/yb/pgsql/TestBatchCopyFrom.java

@@ -27,8 +27,8 @@
 import org.apache.commons.lang3.StringUtils;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.postgresql.copy.CopyManager;
-import org.postgresql.core.BaseConnection;
+import com.yugabyte.copy.CopyManager;
+import com.yugabyte.core.BaseConnection;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.yb.client.TestUtils;

java/yb-pgsql/src/test/java/org/yb/pgsql/TestConsistency.java

@@ -31,8 +31,8 @@
 import org.apache.commons.lang3.time.StopWatch;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.postgresql.util.PSQLException;
-import org.postgresql.util.PSQLState;
+import com.yugabyte.util.PSQLException;
+import com.yugabyte.util.PSQLState;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 

java/yb-pgsql/src/test/java/org/yb/pgsql/TestLoadBalance.java

@@ -14,6 +14,8 @@
 package org.yb.pgsql;
 
 import com.google.common.net.HostAndPort;
+import com.yugabyte.ysql.ClusterAwareLoadBalancer;
+import com.yugabyte.jdbc.PgConnection;
 import org.yb.AssertionWrappers;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -22,15 +24,25 @@
 import org.yb.minicluster.MiniYBDaemon;
 import org.yb.util.YBTestRunnerNonTsanOnly;
 
+import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.Statement;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 @RunWith(value = YBTestRunnerNonTsanOnly.class)
 public class TestLoadBalance extends BasePgSQLTest {
   private static final Logger LOG = LoggerFactory.getLogger(TestPgEncryption.class);
 
+  @Override
+  public ConnectionBuilder getConnectionBuilder() {
+    ConnectionBuilder cb = new ConnectionBuilder(miniCluster);
+    cb.setLoadBalance(true);
+    return cb;
+  }
+
   @Test
   public void testYBServersFunction() throws Exception {
     Statement st = connection.createStatement();
@@ -75,6 +87,82 @@ public void testYBServersFunction() throws Exception {
     }
     AssertionWrappers.assertEquals(
       "expected servers started by minicluster", hostPortsDaemonMap.size(), cnt);
-
+    ClusterAwareLoadBalancer clb = ClusterAwareLoadBalancer.instance();
+    AssertionWrappers.assertNotNull(clb);
+    List<Connection> connList = new ArrayList<>();
+    try {
+      Map<String, Integer> hostToNumConnections = new HashMap<>();
+      for (int i = 0; i < 10; i++) {
+        Connection c = getConnectionBuilder().connect();
+        connList.add(c);
+        String host = ((PgConnection)c).getQueryExecutor().getHostSpec().getHost();
+        Integer numConns = 0;
+        if (hostToNumConnections.containsKey(host)) {
+          numConns = hostToNumConnections.get(host);
+          numConns += 1;
+        } else {
+          numConns = 1;
+        }
+        hostToNumConnections.put(host, numConns);
+      }
+      // Add the first connection host port too
+      String firstHost = ((PgConnection)connection).getQueryExecutor().getHostSpec().getHost();
+      Integer numConns = hostToNumConnections.get(firstHost);
+      hostToNumConnections.put(firstHost, numConns+1);
+      clb.printHostToConnMap();
+      AssertionWrappers.assertEquals(3, hostToNumConnections.size());
+      for (Map.Entry<String, Integer> e : hostToNumConnections.entrySet()) {
+        AssertionWrappers.assertTrue(e.getValue() >= 3);
+      }
+    } finally {
+      for (Connection c : connList) c.close();
+    }
+    // Let's close the first connection as well, so that this connection does not interfere
+    // with the accounting done later in the test when multiple threads try to create the
+    // connections at the same time.
+    connection.close();
+    // Now let's test parallel connection attempts. Even then it should be properly balanced
+    class ConnectionRunnable implements Runnable {
+      volatile Connection conn;
+      volatile Exception ex;
+      @Override
+      public void run() {
+        try {
+          conn = getConnectionBuilder().connect();
+        } catch (Exception e) {
+          ex = e;
+        }
+      }
+    }
+    Thread[] threads = new Thread[10];
+    ConnectionRunnable[] runnables = new ConnectionRunnable[10];
+    for(int i=0; i< 10; i++) {
+      runnables[i] = new ConnectionRunnable();
+      threads[i] = new Thread(runnables[i]);
+    }
+    for(Thread t : threads) {
+      t.start();
+    }
+    for(Thread t : threads) {
+      t.join();
+    }
+    Map<String, Integer> hostToNumConnections = new HashMap<>();
+    for (int i = 0; i < 10; i++) {
+      AssertionWrappers.assertNull(runnables[i].ex);
+      Connection c = runnables[i].conn;
+      String host = ((PgConnection)c).getQueryExecutor().getHostSpec().getHost();
+      Integer numConns;
+      if (hostToNumConnections.containsKey(host)) {
+        numConns = hostToNumConnections.get(host);
+        numConns += 1;
+      } else {
+        numConns = 1;
+      }
+      hostToNumConnections.put(host, numConns);
+      c.close();
+    }
+    for (Map.Entry<String, Integer> e : hostToNumConnections.entrySet()) {
+      AssertionWrappers.assertTrue(e.getValue() >= 3);
+    }
   }
 }

java/yb-pgsql/src/test/java/org/yb/pgsql/TestOneOrTwoAdmins.java

@@ -23,7 +23,7 @@
 import java.sql.Statement;
 import java.sql.SQLException;
 
-import org.postgresql.util.PSQLException;
+import com.yugabyte.util.PSQLException;
 import static org.yb.AssertionWrappers.*;
 
 @RunWith(value=YBTestRunnerNonTsanOnly.class)

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgCacheConsistency.java

@@ -17,7 +17,7 @@
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.postgresql.util.PSQLException;
+import com.yugabyte.util.PSQLException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.yb.minicluster.MiniYBCluster;

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgDropDatabase.java

@@ -20,7 +20,7 @@
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.postgresql.util.PSQLException;
+import com.yugabyte.util.PSQLException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgExplicitLocks.java

@@ -24,7 +24,7 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 
-import org.postgresql.util.PSQLException;
+import com.yugabyte.util.PSQLException;
 import static org.yb.AssertionWrappers.*;
 
 @RunWith(value=YBTestRunnerNonTsanOnly.class)

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgMisc.java

@@ -17,8 +17,8 @@
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.postgresql.util.PSQLException;
-import org.postgresql.util.PSQLWarning;
+import com.yugabyte.util.PSQLException;
+import com.yugabyte.util.PSQLWarning;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.yb.util.YBTestRunnerNonTsanOnly;

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgMultiTouchCacheUsage.java

@@ -20,7 +20,7 @@
 import org.yb.util.YBTestRunnerNonTsanOnly;
 import org.yb.client.YBClient;
 import org.yb.minicluster.MiniYBDaemon;
-import org.postgresql.util.PSQLException;
+import com.yugabyte.util.PSQLException;
 import static org.yb.AssertionWrappers.*;
 
 import java.net.MalformedURLException;

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgRpcForwarding.java

@@ -31,7 +31,7 @@
 
 import com.google.common.net.HostAndPort;
 import org.yb.minicluster.MiniYBDaemon;
-import org.postgresql.util.PSQLException;
+import com.yugabyte.util.PSQLException;
 import static org.yb.AssertionWrappers.*;
 
 @RunWith(value=YBTestRunnerNonTsanOnly.class)