Add ConnectorSplit#getRetainedSizeInBytes method

trinodb · Dec 23, 2021 · a4898a7 · a4898a7
1 parent 4a1624a
commit a4898a7
Show file tree

Hide file tree

Showing 67 changed files with 793 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/connector/CatalogName.java b/core/trino-main/src/main/java/io/trino/connector/CatalogName.java
@@ -15,14 +15,18 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonValue;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.Objects;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static java.util.Objects.requireNonNull;
 
 public final class CatalogName
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(CatalogName.class).instanceSize();
+
     private static final String INFORMATION_SCHEMA_CONNECTOR_PREFIX = "$info_schema@";
     private static final String SYSTEM_TABLES_CONNECTOR_PREFIX = "$system@";
 
@@ -66,6 +70,12 @@ public String toString()
         return catalogName;
     }
 
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE
+                + estimatedSizeOf(catalogName);
+    }
+
     public static boolean isInternalSystemConnector(CatalogName catalogName)
     {
         return catalogName.getCatalogName().startsWith(SYSTEM_TABLES_CONNECTOR_PREFIX) ||

diff --git a/...trino-main/src/main/java/io/trino/connector/informationschema/InformationSchemaSplit.java b/...trino-main/src/main/java/io/trino/connector/informationschema/InformationSchemaSplit.java
@@ -18,15 +18,19 @@
 import com.google.common.collect.ImmutableList;
 import io.trino.spi.HostAddress;
 import io.trino.spi.connector.ConnectorSplit;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.List;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static java.util.Objects.requireNonNull;
 
 public class InformationSchemaSplit
         implements ConnectorSplit
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(InformationSchemaSplit.class).instanceSize();
+
     private final List<HostAddress> addresses;
 
     @JsonCreator
@@ -56,4 +60,11 @@ public Object getInfo()
     {
         return this;
     }
+
+    @Override
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE
+                + estimatedSizeOf(addresses, HostAddress::getRetainedSizeInBytes);
+    }
 }
diff --git a/core/trino-main/src/main/java/io/trino/connector/system/SystemColumnHandle.java b/core/trino-main/src/main/java/io/trino/connector/system/SystemColumnHandle.java
@@ -18,15 +18,19 @@
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
 import io.trino.spi.connector.ConnectorTableMetadata;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.Map;
 
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static java.util.Objects.requireNonNull;
 
 public class SystemColumnHandle
         implements ColumnHandle
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(SystemColumnHandle.class).instanceSize();
+
     private final String columnName;
 
     @JsonCreator
@@ -67,6 +71,12 @@ public String toString()
         return columnName;
     }
 
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE
+                + estimatedSizeOf(columnName);
+    }
+
     public static Map<String, ColumnHandle> toSystemColumnHandles(ConnectorTableMetadata tableMetadata)
     {
         return tableMetadata.getColumns().stream().collect(toImmutableMap(

diff --git a/core/trino-main/src/main/java/io/trino/connector/system/SystemSplit.java b/core/trino-main/src/main/java/io/trino/connector/system/SystemSplit.java
@@ -20,16 +20,20 @@
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorSplit;
 import io.trino.spi.predicate.TupleDomain;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.List;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkArgument;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static java.util.Objects.requireNonNull;
 
 public class SystemSplit
         implements ConnectorSplit
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(SystemSplit.class).instanceSize();
+
     private final List<HostAddress> addresses;
     private final TupleDomain<ColumnHandle> constraint;
 
@@ -74,6 +78,14 @@ public Object getInfo()
         return this;
     }
 
+    @Override
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE
+                + estimatedSizeOf(addresses, HostAddress::getRetainedSizeInBytes)
+                + constraint.getRetainedSizeInBytes(columnHandle -> ((SystemColumnHandle) columnHandle).getRetainedSizeInBytes());
+    }
+
     @Override
     public String toString()
     {

diff --git a/core/trino-main/src/main/java/io/trino/execution/TaskId.java b/core/trino-main/src/main/java/io/trino/execution/TaskId.java
@@ -16,17 +16,21 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonValue;
 import io.trino.spi.QueryId;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.List;
 import java.util.Objects;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static io.trino.spi.QueryId.parseDottedId;
 import static java.lang.Integer.parseInt;
 import static java.util.Objects.requireNonNull;
 
 public class TaskId
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(TaskId.class).instanceSize();
+
     @JsonCreator
     public static TaskId valueOf(String taskId)
     {
@@ -94,4 +98,9 @@ public boolean equals(Object obj)
         TaskId other = (TaskId) obj;
         return Objects.equals(this.fullId, other.fullId);
     }
+
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE + estimatedSizeOf(fullId);
+    }
 }
diff --git a/core/trino-main/src/main/java/io/trino/operator/index/IndexSplit.java b/core/trino-main/src/main/java/io/trino/operator/index/IndexSplit.java
@@ -49,6 +49,13 @@ public Object getInfo()
         return null;
     }
 
+    @Override
+    public long getRetainedSizeInBytes()
+    {
+        // IndexSplit is expected to be short lived and is not expected to be queried for the memory it retains
+        throw new UnsupportedOperationException();
+    }
+
     public RecordSet getKeyRecordSet()
     {
         return keyRecordSet;

diff --git a/core/trino-main/src/main/java/io/trino/split/EmptySplit.java b/core/trino-main/src/main/java/io/trino/split/EmptySplit.java
@@ -19,6 +19,7 @@
 import io.trino.connector.CatalogName;
 import io.trino.spi.HostAddress;
 import io.trino.spi.connector.ConnectorSplit;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.List;
 
@@ -27,6 +28,8 @@
 public class EmptySplit
         implements ConnectorSplit
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(EmptySplit.class).instanceSize();
+
     private final CatalogName catalogName;
 
     @JsonCreator
@@ -54,6 +57,13 @@ public Object getInfo()
         return this;
     }
 
+    @Override
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE
+                + catalogName.getRetainedSizeInBytes();
+    }
+
     @JsonProperty
     public CatalogName getCatalogName()
     {

diff --git a/core/trino-main/src/main/java/io/trino/split/RemoteSplit.java b/core/trino-main/src/main/java/io/trino/split/RemoteSplit.java
@@ -19,16 +19,19 @@
 import io.trino.execution.TaskId;
 import io.trino.spi.HostAddress;
 import io.trino.spi.connector.ConnectorSplit;
+import org.openjdk.jol.info.ClassLayout;
 
-import java.net.URI;
 import java.util.List;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static java.util.Objects.requireNonNull;
 
 public class RemoteSplit
         implements ConnectorSplit
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(RemoteSplit.class).instanceSize();
+
     private final TaskId taskId;
     private final String location;
 
@@ -77,4 +80,12 @@ public String toString()
                 .add("location", location)
                 .toString();
     }
+
+    @Override
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE
+                + taskId.getRetainedSizeInBytes()
+                + estimatedSizeOf(location);
+    }
 }
diff --git a/core/trino-main/src/main/java/io/trino/testing/TestingSplit.java b/core/trino-main/src/main/java/io/trino/testing/TestingSplit.java
@@ -18,12 +18,17 @@
 import com.google.common.collect.ImmutableList;
 import io.trino.spi.HostAddress;
 import io.trino.spi.connector.ConnectorSplit;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.util.List;
 
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
+
 public class TestingSplit
         implements ConnectorSplit
 {
+    private static final int INSTANCE_SIZE = ClassLayout.parseClass(TestingSplit.class).instanceSize();
+
     private static final HostAddress localHost = HostAddress.fromString("127.0.0.1");
 
     private final boolean remotelyAccessible;
@@ -70,4 +75,11 @@ public Object getInfo()
     {
         return this;
     }
+
+    @Override
+    public long getRetainedSizeInBytes()
+    {
+        return INSTANCE_SIZE
+                + estimatedSizeOf(addresses, HostAddress::getRetainedSizeInBytes);
+    }
 }
diff --git a/core/trino-main/src/test/java/io/trino/connector/MockConnector.java b/core/trino-main/src/test/java/io/trino/connector/MockConnector.java
@@ -727,5 +727,11 @@ public Object getInfo()
         {
             return "mock connector split";
         }
+
+        @Override
+        public long getRetainedSizeInBytes()
+        {
+            return 0;
+        }
     }
 }
diff --git a/core/trino-main/src/test/java/io/trino/execution/BenchmarkNodeScheduler.java b/core/trino-main/src/test/java/io/trino/execution/BenchmarkNodeScheduler.java
@@ -53,6 +53,7 @@
 import org.openjdk.jmh.annotations.State;
 import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jol.info.ClassLayout;
 
 import java.net.URI;
 import java.util.ArrayList;
@@ -68,6 +69,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static io.trino.SystemSessionProperties.MAX_UNACKNOWLEDGED_SPLITS_PER_TASK;
 import static java.util.concurrent.Executors.newCachedThreadPool;
 import static java.util.concurrent.Executors.newScheduledThreadPool;
@@ -248,6 +250,8 @@ private static TopologyAwareNodeSelectorConfig getBenchmarkNetworkTopologyConfig
     private static class TestSplitRemote
             implements ConnectorSplit
     {
+        private static final int INSTANCE_SIZE = ClassLayout.parseClass(TestSplitRemote.class).instanceSize();
+
         private final List<HostAddress> hosts;
 
         public TestSplitRemote(int dataHost)
@@ -272,6 +276,13 @@ public Object getInfo()
         {
             return this;
         }
+
+        @Override
+        public long getRetainedSizeInBytes()
+        {
+            return INSTANCE_SIZE
+                    + estimatedSizeOf(hosts, HostAddress::getRetainedSizeInBytes);
+        }
     }
 
     private static HostAddress addressForHost(int host)

diff --git a/core/trino-main/src/test/java/io/trino/execution/TestNodeScheduler.java b/core/trino-main/src/test/java/io/trino/execution/TestNodeScheduler.java
@@ -47,6 +47,7 @@
 import io.trino.sql.planner.plan.PlanNodeId;
 import io.trino.testing.TestingSession;
 import io.trino.util.FinalizerService;
+import org.openjdk.jol.info.ClassLayout;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.DataProvider;
@@ -72,6 +73,7 @@
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
+import static io.airlift.slice.SizeOf.estimatedSizeOf;
 import static io.airlift.testing.Assertions.assertLessThanOrEqual;
 import static io.trino.spi.StandardErrorCode.NO_NODES_AVAILABLE;
 import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy;
@@ -869,6 +871,8 @@ private static Session sessionWithMaxUnacknowledgedSplitsPerTask(int maxUnacknow
     private static class TestSplitLocal
             implements ConnectorSplit
     {
+        private static final int INSTANCE_SIZE = ClassLayout.parseClass(TestSplitLocal.class).instanceSize();
+
         private final HostAddress address;
         private final SplitWeight splitWeight;
 
@@ -912,6 +916,14 @@ public SplitWeight getSplitWeight()
             return splitWeight;
         }
 
+        @Override
+        public long getRetainedSizeInBytes()
+        {
+            return INSTANCE_SIZE
+                    + address.getRetainedSizeInBytes()
+                    + splitWeight.getRetainedSizeInBytes();
+        }
+
         @Override
         public String toString()
         {
@@ -941,11 +953,19 @@ public Object getInfo()
         {
             return this;
         }
+
+        @Override
+        public long getRetainedSizeInBytes()
+        {
+            return 0;
+        }
     }
 
     private static class TestSplitRemote
             implements ConnectorSplit
     {
+        private static final int INSTANCE_SIZE = ClassLayout.parseClass(TestSplitRemote.class).instanceSize();
+
         private final List<HostAddress> hosts;
         private final SplitWeight splitWeight;
 
@@ -992,6 +1012,14 @@ public SplitWeight getSplitWeight()
         {
             return splitWeight;
         }
+
+        @Override
+        public long getRetainedSizeInBytes()
+        {
+            return INSTANCE_SIZE
+                    + estimatedSizeOf(hosts, HostAddress::getRetainedSizeInBytes)
+                    + splitWeight.getRetainedSizeInBytes();
+        }
     }
 
     private static class TestNetworkTopology