[yugabyte#6509] YSQL: Fix leaking when a portal is used to query in s…

…mall batches Summary: **(A) Fix Leaking** The leak happens when `PortalRun` is called many times to run SELECT in small batches. The memory that is private to each batch is allocated from PortalContext, so the PortalContext keeps growing as the number of batches increases. To fix this, portal's `RunContext` is introduced. Currently, each portal has two memory contexts. - `PortalContext` lasts for the lifetime of the Portal. - `TmpContext` lasts for constructing one row in the result set. To fix this leak, `RunContext` is added. This context will live for one `PortalRun()`. It lasts longer than `TmpContext` but shorter than `PortalContext`. **(B) Testing** 1. Setting up "yb-sample" project to run `maven` against local DB to track memory usage. -- This project creates table, inserts data, scans, and reports memory usage at each row. -- Command example: To run `LargeScan.java` program in this project, use the command "//`mvn install exec:java -Dexec.mainClass=org.yb.sample.LargeScan`//" 2. Introduce SQL builtin functions to collect memory usage at a glance. -- `yb_getrusage()` returns the output from system call `getrusage()`. -- `yb_mem_usage_kb()` returns memory usage in kbs by current session in proxy server. -- `yb_mem_usage()` is the text version of `yb_mem_usage_kb()`. Returning text value allows Yugabyte to provide more details than just the total usage. -- `yb_mem_usage_sql_b` and ```yb_mem_usage_sql_kb` returns memory usage in bytes and kbs respectively by SQL layer of the current session in proxy server. -- `yb_mem_usage_sql` return SQL usage in text. The text value maybe contain more information than the total usage. Test Plan: TestPgPortalLeak Reviewers: mihnea Reviewed By: mihnea Subscribers: mbautin, alex, yql Differential Revision: https://phabricator.dev.yugabyte.com/D11331
bozhiyou · May 7, 2021 · b6145a4 · b6145a4
1 parent 577c871
commit b6145a4
Show file tree

Hide file tree

Showing 21 changed files with 933 additions and 3 deletions.
diff --git a/java/pom.xml b/java/pom.xml
@@ -136,6 +136,7 @@
     <module>yb-jedis-tests</module>
     <module>yb-cdc</module>
     <module>yb-multiapi</module>
+    <module>yb-sample</module>
   </modules>
 
   <build>
@@ -361,6 +362,11 @@
         <artifactId>yb-pgsql</artifactId>
         <version>${project.version}</version>
       </dependency>
+      <dependency>
+        <groupId>org.yb</groupId>
+        <artifactId>yb-sample</artifactId>
+        <version>${project.version}</version>
+      </dependency>
       <dependency>
         <groupId>log4j</groupId>
         <artifactId>log4j</artifactId>

diff --git a/java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgFlags.java b/java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgFlags.java
@@ -0,0 +1,48 @@
+// Copyright (c) YugaByte, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied.  See the License for the specific language governing permissions and limitations
+// under the License.
+//
+
+package org.yb.pgsql;
+import org.yb.pgsql.TestPgPortalLeak;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.yb.util.YBTestRunnerNonTsanOnly;
+import static org.yb.AssertionWrappers.*;
+
+import java.util.Map;
+
+import com.google.common.collect.ImmutableMap;
+
+// This module is used to verify that certain GFLAGS are setup correctly. It only sets the flag
+// value, but the associated feature tests should be defined in their own file.
+@RunWith(value=YBTestRunnerNonTsanOnly.class)
+public class TestPgFlags extends BasePgSQLTest {
+  private static final Logger LOG = LoggerFactory.getLogger(TestPgFlags.class);
+
+  @Override
+  protected Map<String, String> getTServerFlags() {
+    Map<String, String> flagMap = super.getTServerFlags();
+
+    // Flag "ysql_disable_portal_run_context" is used for "TestPgPortalLeak.java".
+    flagMap.put("ysql_disable_portal_run_context", "true");
+
+    return flagMap;
+  }
+
+  @Test
+  public void testPgPortalLeakFlag() throws Exception {
+    TestPgPortalLeak.testPgPortalLeakFlag();
+  }
+}
diff --git a/java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgPortalLeak.java b/java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgPortalLeak.java
@@ -0,0 +1,126 @@
+// Copyright (c) YugaByte, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied.  See the License for the specific language governing permissions and limitations
+// under the License.
+//
+
+package org.yb.pgsql;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.yb.util.YBTestRunnerNonTsanOnly;
+import static org.yb.AssertionWrappers.*;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Map;
+
+@RunWith(value=YBTestRunnerNonTsanOnly.class)
+public class TestPgPortalLeak extends BasePgSQLTest {
+  private static final Logger LOG = LoggerFactory.getLogger(TestPgPortalLeak.class);
+
+  private final int rowCount = 4000;
+
+  private void createTable(String tableName) throws Exception {
+    try (Statement statement = connection.createStatement()) {
+      statement.execute(
+        "CREATE TABLE " + tableName +
+        "  (id text, ename text, age int, city text, about_me text, PRIMARY KEY(id, ename))");
+
+      String insertFormat = "INSERT INTO " + tableName + " VALUES ('%s', '%s', %d, '%s', '%s')";
+      for (int iter = 0; iter < rowCount; iter++) {
+        String id = String.format("user-%04096d", iter);
+        String ename = String.format("name-%d", iter);
+        int age = 20 + iter%50;
+        String city = String.format("city-%d", iter%1000);
+        String aboutMe = String.format("about_me-%d", iter);
+
+        statement.execute(String.format(insertFormat, id, ename, age, city, aboutMe));
+      }
+    }
+  }
+
+  private void selectBatches(String tableName, boolean expectNoLeak) throws Exception {
+    // Set auto commit to false so that query can be read in small batches.
+    connection.setAutoCommit(false);
+
+    try (Statement statement = connection.createStatement()) {
+      // Start transaction.
+      statement.execute("BEGIN");
+
+      // Read data from table in small batches.
+      int fetchSize = 100;
+      String selectTxt = "select yb_mem_usage_sql_b(), id, ename, age, city from " + tableName;
+      Statement selectStmt = connection.createStatement();
+
+      selectStmt.setFetchSize(fetchSize);
+      long expectedUsage = 0;
+      long currentUsage = 0;
+      try (ResultSet rs = selectStmt.executeQuery(selectTxt)) {
+        int rowCount = 0;
+        while (rs.next()) {
+          rowCount++;
+
+          // Memory usage for each batch should be the same as other batches, but the usage will
+          // fluctuate for the first few rows in each batch.
+          if (expectNoLeak) {
+            if (rowCount % fetchSize > 2) {
+              if (expectedUsage == 0) {
+                expectedUsage = rs.getLong(1);
+              }
+              currentUsage = rs.getLong(1);
+              assertEquals(expectedUsage, currentUsage);
+            }
+
+          } else if (rowCount % fetchSize == 7) {
+            // Expecting leaking from batch to batch.
+            currentUsage = rs.getLong(1);
+            assertLessThan(expectedUsage, currentUsage);
+            expectedUsage = currentUsage;
+          }
+
+          // Print result every 500.
+          if (rowCount % 500 == 0) {
+            LOG.info(String.format("Row %d: usage = %d bytes," +
+                                   " ename = '%s', age = '%s', city = '%s'",
+                                   rowCount, rs.getLong(1),
+                                   rs.getString(3).trim(), rs.getString(4), rs.getString(5)));
+          }
+        }
+      } catch (Exception e) {
+        statement.execute("ABORT");
+      }
+
+      // Start transaction.
+      statement.execute("END");
+    }
+  }
+
+  // This test case will be run with default setting where there should be no leak.
+  //   "ysql_disable_portal_run_context == true".
+  @Test
+  public void testNoPortalLeak() throws Exception {
+    String tableName = "tableExpectPgPortalHasNoLeak";
+    createTable(tableName);
+    selectBatches(tableName, true /* expectNoLeak */);
+  }
+
+  // This test case will be run in TestPgFlags() module where leaking is expected.
+  //   "ysql_disable_portal_run_context == true".
+  public static void testPgPortalLeakFlag() throws Exception {
+    String tableName = "tableExpectPgPortalHasLeak";
+    TestPgPortalLeak testCase = new TestPgPortalLeak();
+    testCase.createTable(tableName);
+    testCase.selectBatches(tableName, false /* expectNoLeak */);
+  }
+}
diff --git a/java/yb-sample/pom.xml b/java/yb-sample/pom.xml
@@ -0,0 +1,80 @@
+<?xml version="1.0"?>
+<project
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0
+                        http://maven.apache.org/xsd/maven-4.0.0.xsd"
+    xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>org.yb</groupId>
+    <artifactId>yb-parent</artifactId>
+    <version>0.8.3-SNAPSHOT</version>
+  </parent>
+  <artifactId>yb-sample</artifactId>
+  <name>YB Manual Support</name>
+  <url>http://maven.apache.org</url>
+  <dependencies>
+    <dependency>
+      <groupId>org.yb</groupId>
+      <artifactId>yb-client</artifactId>
+      <version>0.8.3-SNAPSHOT</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.yb</groupId>
+      <artifactId>yb-client</artifactId>
+      <version>0.8.3-SNAPSHOT</version>
+      <type>test-jar</type>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.postgresql</groupId>
+      <artifactId>postgresql</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.json</groupId>
+      <artifactId>json</artifactId>
+      <version>20180130</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.yugabyte</groupId>
+      <artifactId>cassandra-driver-core</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>net.jpountz.lz4</groupId>
+      <artifactId>lz4</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.xerial.snappy</groupId>
+      <artifactId>snappy-java</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>log4j</groupId>
+      <artifactId>log4j</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.slf4j</groupId>
+      <artifactId>slf4j-log4j12</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.slf4j</groupId>
+      <artifactId>slf4j-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>commons-io</groupId>
+      <artifactId>commons-io</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-lang3</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>${junit.groupId}</groupId>
+      <artifactId>junit</artifactId>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+</project>
diff --git a/java/yb-sample/src/main/java/org/yb/LargeScan.java b/java/yb-sample/src/main/java/org/yb/LargeScan.java
@@ -0,0 +1,101 @@
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied.  See the License for the specific language governing permissions and limitations
+// under the License.
+//
+package org.yb.sample;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.*;
+import java.lang.InterruptedException;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+// LargeScan.
+// This module insert a larger number of rows and query them in smaller batches.
+// To install and execute
+//   mvn install exec:java -Dexec.mainClass=org.yb.sample.LargeScan
+public class LargeScan {
+  private static final Logger LOG = LoggerFactory.getLogger(LargeScan.class);
+
+  private static int rowCount = 4000;
+
+  private static void createTableUsers(Statement stmt, YbSqlUtil yb) throws Exception {
+    yb.exec(stmt,
+            "CREATE TABLE IF NOT EXISTS users" +
+            "  (id text, ename text, age int, city text, about_me text, PRIMARY KEY(id, ename))");
+
+    String insertFormat = "INSERT INTO users VALUES ('%s', '%s', %d, '%s', '%s')";
+    for (int iter = 0; iter < rowCount; iter++) {
+      String id = String.format("user-%04096d", iter);
+      String ename = String.format("name-%d", iter);
+      int age = 20 + iter%50;
+      String city = String.format("city-%d", iter%1000);
+      String aboutMe = String.format("about_me-%d", iter);
+
+      yb.exec(stmt, String.format(insertFormat, id, ename, age, city, aboutMe));
+    }
+  }
+
+  public static void main(String[] args) throws Exception {
+    YbSqlUtil yb = new YbSqlUtil();
+
+    // Connect to local YB database.
+    Connection cxn = yb.connectLocal();
+
+    // Get the process id for debugging purpose.
+    ResultSet pidResult = yb.execQuery("SELECT pg_backend_pid()");
+    pidResult.next();
+    LOG.info(String.format("SELECT process ID = %d", pidResult.getInt(1)));
+    pidResult.close();
+
+    try {
+      // Setup large table "users" if needed.
+      if (!yb.tableExists("users")) {
+        Statement createStmt = cxn.createStatement();
+        createTableUsers(createStmt, yb);
+      }
+
+      // Start transaction.
+      cxn.setAutoCommit(false);
+      yb.exec("BEGIN");
+
+      // Read data from table.
+      try (Statement selectStmt = cxn.createStatement()) {
+        selectStmt.setFetchSize(100);
+        ResultSet rs = yb.execQuery(selectStmt,
+            "select yb_mem_usage_sql_b(), id, ename, age, city from users");
+
+        int rowCount = 0;
+        while (rs.next()) {
+          // Print result every 500.
+          rowCount++;
+          if (rowCount % 500 == 0) {
+            LOG.info(String.format("Row %d: usage = %d bytes," +
+                                   " ename = '%s', age = '%s', city = '%s'",
+                                   rowCount, rs.getLong(1),
+                                   rs.getString(3).trim(), rs.getString(4), rs.getString(5)));
+          }
+        }
+        rs.close();
+      }
+
+      // Close transaction.
+      yb.exec("END");
+
+    } catch (Exception e) {
+      yb.exec("ABORT");
+      LOG.info("Failed to execute LargeScan. " + e.getMessage());
+    }
+  }
+}