[Pipe] Add SHOW PIPE MEMORY for named memory blocks

integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowPipeMemoryIT.java

@@ -0,0 +1,246 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.apache.iotdb.db.it;
+
+import org.apache.iotdb.isession.SessionConfig;
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.env.MultiEnvFactory;
+import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.ClusterIT;
+import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+import org.apache.iotdb.itbase.category.TableClusterIT;
+import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
+import org.apache.iotdb.itbase.env.BaseEnv;
+
+import org.awaitility.Awaitility;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({
+  LocalStandaloneIT.class,
+  ClusterIT.class,
+  TableLocalStandaloneIT.class,
+  TableClusterIT.class
+})
+public class IoTDBShowPipeMemoryIT {
+
+  private static final String PIPE_NAME = "memory_owner";
+  private static final String USER = "pipe_memory_user";
+  private static final String PASSWORD = "password123456";
+  private static final String[] COLUMN_NAMES = {
+    "block_id",
+    "name",
+    "category",
+    "pipe_name",
+    "pipe_model",
+    "memory_usage_in_bytes",
+    "peak_memory_size_in_bytes",
+    "max_memory_size_in_bytes",
+    "allocation_time",
+    "assigner",
+    "parent_block_id",
+    "hierarchy_level",
+    "accounted_memory_usage_in_bytes"
+  };
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false);
+    // Reserve enough pipe memory for both pipes and the shared caches in the small test JVM.
+    EnvFactory.getEnv().getConfig().getDataNodeConfig().setDataNodeMemoryProportion("3:3:1:1:9:0");
+    EnvFactory.getEnv().initClusterEnvironment();
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("CREATE USER " + USER + " '" + PASSWORD + "'");
+    }
+  }
+
+  @AfterClass
+  public static void tearDown() throws Exception {
+    EnvFactory.getEnv().cleanClusterEnvironment();
+  }
+
+  @Test
+  public void testColumnsAndPrivilegesInBothModels() throws Exception {
+    for (DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) {
+      for (String dialect : Arrays.asList(BaseEnv.TREE_SQL_DIALECT, BaseEnv.TABLE_SQL_DIALECT)) {
+        try (Connection connection =
+                EnvFactory.getEnv()
+                    .getConnection(
+                        dataNode,
+                        SessionConfig.DEFAULT_USER,
+                        SessionConfig.DEFAULT_PASSWORD,
+                        dialect);
+            Statement statement = connection.createStatement()) {
+          assertPipeMemoryColumnsAndReadOwners(statement, "SHOW PIPE MEMORY");
+        }
+        try (Connection connection =
+                EnvFactory.getEnv().getConnection(dataNode, USER, PASSWORD, dialect);
+            Statement statement = connection.createStatement()) {
+          assertThrows(SQLException.class, () -> statement.executeQuery("SHOW PIPE MEMORY"));
+        }
+      }
+    }
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("GRANT SYSTEM ON root.** TO USER " + USER);
+      try {
+        for (DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) {
+          for (String dialect :
+              Arrays.asList(BaseEnv.TREE_SQL_DIALECT, BaseEnv.TABLE_SQL_DIALECT)) {
+            try (Connection userConnection =
+                    EnvFactory.getEnv().getConnection(dataNode, USER, PASSWORD, dialect);
+                Statement userStatement = userConnection.createStatement()) {
+              assertPipeMemoryColumnsAndReadOwners(userStatement, "SHOW PIPE MEMORY");
+            }
+          }
+        }
+      } finally {
+        statement.execute("REVOKE SYSTEM ON root.** FROM USER " + USER);
+      }
+    }
+  }
+
+  @Test
+  public void testTreeAndTablePipeOwnership() throws Exception {
+    MultiEnvFactory.createEnv(1);
+    final BaseEnv receiver = MultiEnvFactory.getEnv(0);
+    receiver.getConfig().getCommonConfig().setIsPipeEnableMemoryCheck(false);
+    receiver.initClusterEnvironment(1, 1);
+    try (Connection treeConnection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
+        Statement treeStatement = treeConnection.createStatement();
+        Connection tableConnection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement tableStatement = tableConnection.createStatement()) {
+      treeStatement.execute("CREATE DATABASE root.pipe_memory");
+      treeStatement.execute("CREATE TIMESERIES root.pipe_memory.d.s INT32");
+      treeStatement.execute("INSERT INTO root.pipe_memory.d(time, s) VALUES (0, 0)");
+      tableStatement.execute("CREATE DATABASE pipe_memory_db");
+      tableStatement.execute("CREATE TABLE pipe_memory_db.t (device STRING TAG, s INT32 FIELD)");
+      tableStatement.execute("INSERT INTO pipe_memory_db.t(time, device, s) VALUES (0, 'd', 0)");
+      final String createPipe =
+          "CREATE PIPE "
+              + PIPE_NAME
+              + " WITH SOURCE ('source.realtime.mode'='forced-log', 'source.history.enable'='false')"
+              + " WITH SINK ('sink'='iotdb-thrift-sync-sink', 'sink.ip'='"
+              + receiver.getDataNodeWrapper(0).getIp()
+              + "', 'sink.port'='"
+              + receiver.getDataNodeWrapper(0).getPort()
+              + "',"
+              + " 'sink.batch.enable'='false')";
+      try {
+        treeStatement.execute(createPipe);
+        tableStatement.execute(createPipe);
+        // Complete the handshake before stopping the receiver, so each pipe retains its events.
+        receiver.shutdownDataNode(0);
+        treeStatement.execute("INSERT INTO root.pipe_memory.d(time, s) VALUES (1, 1)");
+        tableStatement.execute("INSERT INTO pipe_memory_db.t(time, device, s) VALUES (1, 'd', 1)");
+        final Set<String> expectedOwners = new HashSet<>(Arrays.asList("tree", "table"));
+        Awaitility.await()
+            .atMost(60, TimeUnit.SECONDS)
+            .untilAsserted(
+                () -> {
+                  assertTrue(
+                      readPipeOwnersOnAllDataNodes(BaseEnv.TREE_SQL_DIALECT, "SHOW PIPE MEMORY")
+                          .containsAll(expectedOwners));
+                  assertTrue(
+                      readPipeOwnersOnAllDataNodes(BaseEnv.TABLE_SQL_DIALECT, "SHOW PIPE MEMORY")
+                          .containsAll(expectedOwners));
+                  assertTrue(
+                      readPipeOwnersOnAllDataNodes(
+                              BaseEnv.TABLE_SQL_DIALECT,
+                              "SELECT * FROM information_schema.pipe_memory")
+                          .containsAll(expectedOwners));
+                });
+      } finally {
+        treeStatement.execute("DROP PIPE IF EXISTS " + PIPE_NAME);
+        tableStatement.execute("DROP PIPE IF EXISTS " + PIPE_NAME);
+      }
+    } finally {
+      receiver.cleanClusterEnvironment();
+    }
+  }
+
+  private static Set<String> readPipeOwnersOnAllDataNodes(String dialect, String sql)
+      throws SQLException {
+    final Set<String> owners = new HashSet<>();
+    // Query each node separately because memory diagnostics describe local allocations.
+    for (DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) {
+      try (Connection connection =
+              EnvFactory.getEnv()
+                  .getConnection(
+                      dataNode,
+                      SessionConfig.DEFAULT_USER,
+                      SessionConfig.DEFAULT_PASSWORD,
+                      dialect);
+          Statement statement = connection.createStatement()) {
+        owners.addAll(assertPipeMemoryColumnsAndReadOwners(statement, sql));
+      }
+    }
+    return owners;
+  }
+
+  private static Set<String> assertPipeMemoryColumnsAndReadOwners(Statement statement, String sql)
+      throws SQLException {
+    final Set<String> owners = new HashSet<>();
+    try (ResultSet resultSet = statement.executeQuery(sql)) {
+      final ResultSetMetaData metadata = resultSet.getMetaData();
+      assertEquals(COLUMN_NAMES.length, metadata.getColumnCount());
+      for (int i = 0; i < COLUMN_NAMES.length; i++) {
+        assertEquals(COLUMN_NAMES[i], metadata.getColumnName(i + 1));
+      }
+      boolean hasFloatingMemory = false;
+      while (resultSet.next()) {
+        if ("FloatingMemory".equals(resultSet.getString("name"))) {
+          assertNull(resultSet.getObject("pipe_name"));
+          assertNull(resultSet.getObject("pipe_model"));
+          assertTrue(resultSet.getLong("memory_usage_in_bytes") >= 0);
+          hasFloatingMemory = true;
+        }
+        if (PIPE_NAME.equals(resultSet.getString("pipe_name"))) {
+          final String model = resultSet.getString("pipe_model");
+          assertTrue("tree".equals(model) || "table".equals(model));
+          owners.add(model);
+        }
+      }
+      assertTrue(hasFloatingMemory);
+    }
+    return owners;
+  }
+}

integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBDatabaseIT.java

@@ -536,6 +536,7 @@ public void testInformationSchema() throws SQLException {
               "functions,INF,",
               "keywords,INF,",
               "nodes,INF,",
+              "pipe_memory,INF,",
               "pipe_plugins,INF,",
               "pipes,INF,",
               "queries,INF,",
@@ -618,6 +619,24 @@ public void testInformationSchema() throws SQLException {
                   "estimated_remaining_seconds,DOUBLE,ATTRIBUTE,",
                   "is_degraded,BOOLEAN,ATTRIBUTE,",
                   "recent_failures,STRING,ATTRIBUTE,")));
+      TestUtils.assertResultSetEqual(
+          statement.executeQuery("desc pipe_memory"),
+          "ColumnName,DataType,Category,",
+          new HashSet<>(
+              Arrays.asList(
+                  "block_id,INT64,TAG,",
+                  "name,STRING,TAG,",
+                  "category,STRING,TAG,",
+                  "pipe_name,STRING,ATTRIBUTE,",
+                  "pipe_model,STRING,ATTRIBUTE,",
+                  "memory_usage_in_bytes,INT64,ATTRIBUTE,",
+                  "peak_memory_size_in_bytes,INT64,ATTRIBUTE,",
+                  "max_memory_size_in_bytes,INT64,ATTRIBUTE,",
+                  "allocation_time,TIMESTAMP,ATTRIBUTE,",
+                  "assigner,STRING,ATTRIBUTE,",
+                  "parent_block_id,INT64,ATTRIBUTE,",
+                  "hierarchy_level,INT32,ATTRIBUTE,",
+                  "accounted_memory_usage_in_bytes,INT64,ATTRIBUTE,")));
       TestUtils.assertResultSetEqual(
           statement.executeQuery("desc receivers"),
           "ColumnName,DataType,Category,",
@@ -751,6 +770,9 @@ public void testInformationSchema() throws SQLException {
       Assert.assertThrows(SQLException.class, () -> statement.execute("select * from data_nodes"));
       Assert.assertThrows(
           SQLException.class, () -> statement.executeQuery("select * from pipe_plugins"));
+      Assert.assertThrows(
+          SQLException.class, () -> statement.executeQuery("select * from pipe_memory"));
+      Assert.assertThrows(SQLException.class, () -> statement.executeQuery("SHOW PIPE MEMORY"));
       Assert.assertThrows(
           SQLException.class, () -> statement.executeQuery("select * from table_disk_usage"));
 
@@ -785,6 +807,47 @@ public void testInformationSchema() throws SQLException {
       // Test table query
       statement.execute("use information_schema");
 
+      try (final ResultSet resultSet = statement.executeQuery("SHOW PIPE MEMORY")) {
+        final ResultSetMetaData metaData = resultSet.getMetaData();
+        assertEquals(13, metaData.getColumnCount());
+        assertEquals("block_id", metaData.getColumnName(1));
+        assertEquals("name", metaData.getColumnName(2));
+        assertEquals("category", metaData.getColumnName(3));
+        assertEquals("pipe_name", metaData.getColumnName(4));
+        assertEquals("pipe_model", metaData.getColumnName(5));
+        assertEquals("memory_usage_in_bytes", metaData.getColumnName(6));
+        assertEquals("peak_memory_size_in_bytes", metaData.getColumnName(7));
+        assertEquals("max_memory_size_in_bytes", metaData.getColumnName(8));
+        assertEquals("allocation_time", metaData.getColumnName(9));
+        assertEquals("assigner", metaData.getColumnName(10));
+        assertEquals("parent_block_id", metaData.getColumnName(11));
+        assertEquals("hierarchy_level", metaData.getColumnName(12));
+        assertEquals("accounted_memory_usage_in_bytes", metaData.getColumnName(13));
+        boolean hasFloatingMemory = false;
+        while (resultSet.next()) {
+          if ("FloatingMemory".equals(resultSet.getString(2))) {
+            assertTrue(resultSet.getLong("memory_usage_in_bytes") >= 0);
+            assertNull(resultSet.getObject("pipe_name"));
+            assertNull(resultSet.getObject("pipe_model"));
+            hasFloatingMemory = true;
+          }
+        }
+        assertTrue(hasFloatingMemory);
+      }
+      try (final ResultSet resultSet =
+          statement.executeQuery("select * from information_schema.pipe_memory")) {
+        boolean hasFloatingMemory = false;
+        while (resultSet.next()) {
+          if ("FloatingMemory".equals(resultSet.getString(2))) {
+            assertTrue(resultSet.getLong("memory_usage_in_bytes") >= 0);
+            assertNull(resultSet.getObject("pipe_name"));
+            assertNull(resultSet.getObject("pipe_model"));
+            hasFloatingMemory = true;
+          }
+        }
+        assertTrue(hasFloatingMemory);
+      }
+
       statement.execute("create database test");
       statement.execute(
           "create table test.test (a tag, b attribute, c int32 comment 'turbine') comment 'test'");
@@ -834,6 +897,7 @@ public void testInformationSchema() throws SQLException {
                   "information_schema,columns,INF,USING,null,SYSTEM VIEW,false,",
                   "information_schema,queries,INF,USING,null,SYSTEM VIEW,false,",
                   "information_schema,regions,INF,USING,null,SYSTEM VIEW,false,",
+                  "information_schema,pipe_memory,INF,USING,null,SYSTEM VIEW,false,",
                   "information_schema,topics,INF,USING,null,SYSTEM VIEW,false,",
                   "information_schema,pipe_plugins,INF,USING,null,SYSTEM VIEW,false,",
                   "information_schema,pipes,INF,USING,null,SYSTEM VIEW,false,",
@@ -853,10 +917,11 @@ public void testInformationSchema() throws SQLException {
                   "information_schema,queries_costs_histogram,INF,USING,null,SYSTEM VIEW,false,",
                   "test,test,INF,USING,test,BASE TABLE,true,",
                   "test,view_table,100,USING,null,VIEW FROM TREE,true,")));
+      // Include all system tables and the table and view created in the test database.
       TestUtils.assertResultSetEqual(
           statement.executeQuery("count devices from tables where status = 'USING'"),
           "count(devices),",
-          Collections.singleton("25,"));
+          Collections.singleton((InformationSchema.getSchemaTables().size() + 2) + ","));
       TestUtils.assertResultSetEqual(
           statement.executeQuery(
               "select * from columns where table_name = 'queries' or database = 'test'"),

iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4

@@ -148,6 +148,7 @@ keyWords
     | LOAD
     | LOCAL
     | LOCK
+    | MEMORY
     | MERGE
     | METADATA
     | MIGRATE

iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4

@@ -55,7 +55,7 @@ ddlStatement
     // ExternalService
     | createService | startService | stopService | dropService | showService
     // Pipe Task
-    | createPipe | alterPipe | dropPipe | startPipe | stopPipe | showPipes | showReceivers
+    | createPipe | alterPipe | dropPipe | startPipe | stopPipe | showPipeMemory | showPipes | showReceivers
     // Pipe Plugin
     | createPipePlugin | dropPipePlugin | showPipePlugins
     // Subscription
@@ -702,6 +702,10 @@ showPipes
     : SHOW ((PIPE pipeName=identifier) | PIPES (WHERE (CONNECTOR | SINK) USED BY pipeName=identifier)?)
     ;
 
+showPipeMemory
+    : SHOW PIPE MEMORY
+    ;
+
 showReceivers
     : SHOW RECEIVERS
     ;

iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4

@@ -523,6 +523,10 @@ LOCK
     : L O C K
     ;
 
+MEMORY
+    : M E M O R Y
+    ;
+
 MERGE
     : M E R G E
     ;