Bound and measure BigQuery's Arrow buffer allocations

plugin/trino-bigquery/pom.xml

@@ -223,6 +223,11 @@
             <artifactId>log</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>stats</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>units</artifactId>
@@ -323,6 +328,11 @@
             <artifactId>threetenbp</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>org.weakref</groupId>
+            <artifactId>jmxutils</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-annotations</artifactId>

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowAllocatorStats.java

@@ -0,0 +1,82 @@
+/*
+ * 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 io.trino.plugin.bigquery;
+
+import io.airlift.stats.CounterStat;
+import org.apache.arrow.memory.AllocationListener;
+import org.apache.arrow.memory.AllocationOutcome;
+import org.weakref.jmx.Managed;
+import org.weakref.jmx.Nested;
+
+public class BigQueryArrowAllocatorStats
+        implements AllocationListener
+{
+    private final CounterStat preAllocatedMemory = new CounterStat();
+    private final CounterStat allocatedMemory = new CounterStat();
+    private final CounterStat releasedMemory = new CounterStat();
+    private final CounterStat failedMemory = new CounterStat();
+
+    @Managed
+    @Nested
+    public CounterStat getAllocatedMemory()
+    {
+        return allocatedMemory;
+    }
+
+    @Managed
+    @Nested
+    public CounterStat getReleasedMemory()
+    {
+        return releasedMemory;
+    }
+
+    @Managed
+    @Nested
+    public CounterStat getPreAllocatedMemory()
+    {
+        return preAllocatedMemory;
+    }
+
+    @Managed
+    @Nested
+    public CounterStat getFailedMemory()
+    {
+        return failedMemory;
+    }
+
+    @Override
+    public void onAllocation(long size)
+    {
+        allocatedMemory.update(size);
+    }
+
+    @Override
+    public void onRelease(long size)
+    {
+        releasedMemory.update(size);
+    }
+
+    @Override
+    public void onPreAllocation(long size)
+    {
+        preAllocatedMemory.update(size);
+    }
+
+    @Override
+    public boolean onFailedAllocation(long size, AllocationOutcome outcome)
+    {
+        failedMemory.update(size);
+        return false;
+    }
+}

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowBufferAllocator.java

@@ -0,0 +1,125 @@
+/*
+ * 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 io.trino.plugin.bigquery;
+
+import com.google.inject.Inject;
+import io.airlift.log.Logger;
+import org.apache.arrow.memory.AllocationListener;
+import org.apache.arrow.memory.AllocationOutcome;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.weakref.jmx.Flatten;
+import org.weakref.jmx.Managed;
+
+import static java.util.Objects.requireNonNull;
+
+public class BigQueryArrowBufferAllocator
+{
+    private static final Logger log = Logger.get(BigQueryArrowBufferAllocator.class);
+
+    private final long maximumAllocation;
+    private final BigQueryArrowAllocatorStats stats;
+    private final BufferAllocator rootAllocator;
+
+    @Inject
+    public BigQueryArrowBufferAllocator(BigQueryArrowConfig config, BigQueryArrowAllocatorStats stats)
+    {
+        this.maximumAllocation = requireNonNull(config, "config is null").getMaxAllocation().toBytes();
+        this.stats = requireNonNull(stats, "stats is null");
+        this.rootAllocator = new RootAllocator(stats, maximumAllocation);
+    }
+
+    public BufferAllocator newChildAllocator(BigQuerySplit split)
+    {
+        return rootAllocator.newChildAllocator(
+                split.streamName(),
+                new RetryingAllocationListener(split.streamName(), stats),
+                split.dataSize().orElse(0),
+                maximumAllocation);
+    }
+
+    @Managed
+    @Flatten
+    public BigQueryArrowAllocatorStats getStats()
+    {
+        return stats;
+    }
+
+    @Managed
+    public long getCurrentAllocatedMemory()
+    {
+        return rootAllocator.getAllocatedMemory();
+    }
+
+    @Managed
+    public long getPeakAllocatedMemory()
+    {
+        return rootAllocator.getPeakMemoryAllocation();
+    }
+
+    @Managed
+    public long getCurrentMemoryHeadroom()
+    {
+        return rootAllocator.getHeadroom();
+    }
+
+    @Managed
+    public long getChildAllocatorsCount()
+    {
+        return rootAllocator.getChildAllocators().size();
+    }
+
+    private static class RetryingAllocationListener
+            implements AllocationListener
+    {
+        private final String name;
+        private final BigQueryArrowAllocatorStats stats;
+
+        private RetryingAllocationListener(String name, BigQueryArrowAllocatorStats stats)
+        {
+            this.name = requireNonNull(name, "name is null");
+            this.stats = requireNonNull(stats, "stats is null");
+        }
+
+        @Override
+        public boolean onFailedAllocation(long size, AllocationOutcome outcome)
+        {
+            stats.onFailedAllocation(size, outcome);
+
+            log.warn("Failed to allocate %d bytes for allocator '%s' due to %s", size, name, outcome.getStatus());
+            outcome.getDetails().ifPresent(details -> {
+                log.warn("Allocation failure details: %s", details.toString());
+            });
+            return false;
+        }
+
+        @Override
+        public void onPreAllocation(long size)
+        {
+            stats.onPreAllocation(size);
+        }
+
+        @Override
+        public void onAllocation(long size)
+        {
+            stats.onAllocation(size);
+        }
+
+        @Override
+        public void onRelease(long size)
+        {
+            stats.onRelease(size);
+        }
+    }
+}

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowConfig.java

@@ -0,0 +1,36 @@
+/*
+ * 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 io.trino.plugin.bigquery;
+
+import io.airlift.configuration.Config;
+import io.airlift.configuration.ConfigDescription;
+import io.airlift.units.DataSize;
+
+public class BigQueryArrowConfig
+{
+    private DataSize maxAllocation = DataSize.ofBytes(Integer.MAX_VALUE);
+
+    public DataSize getMaxAllocation()
+    {
+        return maxAllocation;
+    }
+
+    @ConfigDescription("Maximum memory allocation allowed for Arrow library")
+    @Config("bigquery.arrow-serialization.max-allocation")
+    public BigQueryArrowConfig setMaxAllocation(DataSize maxAllocation)
+    {
+        this.maxAllocation = maxAllocation;
+        return this;
+    }
+}

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryConnectorFactory.java

@@ -18,12 +18,15 @@
 import io.airlift.json.JsonModule;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Tracer;
+import io.trino.plugin.base.jmx.ConnectorObjectNameGeneratorModule;
+import io.trino.plugin.base.jmx.MBeanServerModule;
 import io.trino.spi.NodeManager;
 import io.trino.spi.catalog.CatalogName;
 import io.trino.spi.connector.Connector;
 import io.trino.spi.connector.ConnectorContext;
 import io.trino.spi.connector.ConnectorFactory;
 import io.trino.spi.type.TypeManager;
+import org.weakref.jmx.guice.MBeanModule;
 
 import java.util.Map;
 
@@ -49,6 +52,9 @@ public Connector create(String catalogName, Map<String, String> config, Connecto
         Bootstrap app = new Bootstrap(
                 new JsonModule(),
                 new BigQueryConnectorModule(),
+                new MBeanServerModule(),
+                new MBeanModule(),
+                new ConnectorObjectNameGeneratorModule("io.trino.plugin.bigquery", "trino.plugin.bigquery"),
                 binder -> {
                     binder.bind(TypeManager.class).toInstance(context.getTypeManager());
                     binder.bind(NodeManager.class).toInstance(context.getNodeManager());