Repository: polarsignals/frostdb
Pull request #477 (2023-07-25 13:51:18)

Limit allocation fixes for panic
This adds fixes for leaking allocation when the limit allocator panics.

Of note it adds the `Close()` method to output plans that is defer called during Execute to ensure resources are Released.
It also fixes a couple of other corner cases of leaking resources on panic.


Patch:
diff --git a/go.mod b/go.mod
index 2a44633..e5a2b35 100644
--- a/go.mod
+++ b/go.mod
@@ -4,7 +4,7 @@ go 1.19
 
 require (
 	github.com/RoaringBitmap/roaring v0.9.4
-	github.com/apache/arrow/go/v13 v13.0.0-20230524164752-c4ea194c6880
+	github.com/apache/arrow/go/v13 v13.0.0-20230725201451-5a9240fd28e1
 	github.com/cespare/xxhash/v2 v2.2.0
 	github.com/cockroachdb/datadriven v1.0.2
 	github.com/dgryski/go-metro v0.0.0-20211217172704-adc40b04c140
@@ -23,14 +23,13 @@ require (
 	go.opentelemetry.io/otel v1.16.0
 	go.opentelemetry.io/otel/trace v1.16.0
 	go.uber.org/goleak v1.1.12
-	golang.org/x/exp v0.0.0-20220827204233-334a2380cb91
+	golang.org/x/exp v0.0.0-20230206171751-46f607a40771
 	golang.org/x/sync v0.3.0
 	google.golang.org/protobuf v1.31.0
 )
 
 require (
 	github.com/andybalholm/brotli v1.0.5 // indirect
-	github.com/apache/thrift v0.16.0 // indirect
 	github.com/benbjohnson/immutable v0.4.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bits-and-blooms/bitset v1.2.0 // indirect
@@ -41,18 +40,14 @@ require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/efficientgo/core v1.0.0-rc.2 // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
-	github.com/goccy/go-json v0.9.11 // indirect
+	github.com/goccy/go-json v0.10.0 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
-	github.com/golang/snappy v0.0.4 // indirect
-	github.com/google/flatbuffers v2.0.8+incompatible // indirect
+	github.com/google/flatbuffers v23.1.21+incompatible // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
-	github.com/klauspost/asmfmt v1.3.2 // indirect
 	github.com/klauspost/compress v1.16.7 // indirect
-	github.com/klauspost/cpuid/v2 v2.1.0 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.3 // indirect
 	github.com/mattn/go-runewidth v0.0.9 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
-	github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 // indirect
-	github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 // indirect
 	github.com/mschoch/smat v0.2.0 // indirect
 	github.com/pierrec/lz4/v4 v4.1.18 // indirect
 	github.com/pingcap/errors v0.11.5-0.20210425183316-da1aaba5fb63 // indirect
@@ -73,7 +68,7 @@ require (
 	golang.org/x/sys v0.10.0 // indirect
 	golang.org/x/text v0.9.0 // indirect
 	golang.org/x/tools v0.6.0 // indirect
-	golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f // indirect
+	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
diff --git a/go.sum b/go.sum
index 982da06..e41bb01 100644
--- a/go.sum
+++ b/go.sum
@@ -1,14 +1,11 @@
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
-github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c h1:RGWPOewvKIROun94nF7v2cua9qP+thov/7M50KEoeSU=
 github.com/RoaringBitmap/roaring v0.9.4 h1:ckvZSX5gwCRaJYBNe7syNawCU5oruY9gQmjXlp4riwo=
 github.com/RoaringBitmap/roaring v0.9.4/go.mod h1:icnadbWcNyfEHlYdr+tDlOTih1Bf/h+rzPpv4sbomAA=
 github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs=
 github.com/andybalholm/brotli v1.0.5/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
-github.com/apache/arrow/go/v13 v13.0.0-20230524164752-c4ea194c6880 h1:dUnLG5Kz1kwDZmAst//PkBiy/p6IWXbBP9d3Q5Ex9XI=
-github.com/apache/arrow/go/v13 v13.0.0-20230524164752-c4ea194c6880/go.mod h1:/XatdE3kDIBqZKhZ7OBUHwP2jaASDFZHqF4puOWM8po=
-github.com/apache/thrift v0.16.0 h1:qEy6UW60iVOlUy+b9ZR0d5WzUWYGOo4HfopoyBaNmoY=
-github.com/apache/thrift v0.16.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2FXSqgU=
+github.com/apache/arrow/go/v13 v13.0.0-20230725201451-5a9240fd28e1 h1:fRZe9ew77QcD7d00X1ZgvAreRJF3yzBflEUEMBv2+rM=
+github.com/apache/arrow/go/v13 v13.0.0-20230725201451-5a9240fd28e1/go.mod h1:W69eByFNO0ZR30q1/7Sr9d83zcVZmF2MiP3fFYAWJOc=
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/benbjohnson/immutable v0.4.0 h1:CTqXbEerYso8YzVPxmWxh2gnoRQbbB9X1quUC8+vGZA=
@@ -43,19 +40,16 @@ github.com/go-kit/log v0.2.1/go.mod h1:NwTd00d/i8cPZ3xOwwiv2PO5MOcx78fFErGNcVmBj
 github.com/go-logfmt/logfmt v0.6.0 h1:wGYYu3uicYdqXVgoYbvnkrPVXkuLM1p1ifugDMEdRi4=
 github.com/go-logfmt/logfmt v0.6.0/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KEVveWlfTs=
 github.com/go-sql-driver/mysql v1.6.0 h1:BCTh4TKNUYmOmMUcQ3IipzF5prigylS7XXjEkfCHuOE=
-github.com/goccy/go-json v0.9.11 h1:/pAaQDLHEoCq/5FFmSKBswWmK6H0e8g4159Kc/X/nqk=
-github.com/goccy/go-json v0.9.11/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
-github.com/golang/mock v1.5.0/go.mod h1:CWnOUgYIOo4TcNZ0wHX3YZCqsaM1I1Jvs6v3mP3KVu8=
+github.com/goccy/go-json v0.10.0 h1:mXKd9Qw4NuzShiRlOXKews24ufknHO7gx30lsDyokKA=
+github.com/goccy/go-json v0.10.0/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
 github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
-github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
-github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
 github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
-github.com/google/flatbuffers v2.0.8+incompatible h1:ivUb1cGomAB101ZM1T0nOiWz9pSrTMoa9+EiY7igmkM=
-github.com/google/flatbuffers v2.0.8+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8=
+github.com/google/flatbuffers v23.1.21+incompatible h1:bUqzx/MXCDxuS0hRJL2EfjyZL3uQrPbMocUa8zGqsTA=
+github.com/google/flatbuffers v23.1.21+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
@@ -63,12 +57,10 @@ github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
-github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4=
-github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
 github.com/klauspost/compress v1.16.7 h1:2mk3MPGNzKyxErAw8YaohYh69+pa4sIQSC0fPGCFR9I=
 github.com/klauspost/compress v1.16.7/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
-github.com/klauspost/cpuid/v2 v2.1.0 h1:eyi1Ad2aNJMW95zcSbmGg7Cg6cq3ADwLpMAP96d8rF0=
-github.com/klauspost/cpuid/v2 v2.1.0/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
+github.com/klauspost/cpuid/v2 v2.2.3 h1:sxCkb+qR91z4vsqw4vGGZlDgPz3G7gjaLyK3V8y70BU=
+github.com/klauspost/cpuid/v2 v2.2.3/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
@@ -78,10 +70,6 @@ github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/Qd
 github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
 github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo=
 github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4=
-github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=
-github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
-github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI=
-github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=
 github.com/mschoch/smat v0.2.0 h1:8imxQsjDm8yFEAVBe7azKmKSgzSkZXDuKkSq9374khM=
 github.com/mschoch/smat v0.2.0/go.mod h1:kc9mz7DoBKqDyiRL7VZN8KvXQMWeTaVnttLRXOlotKw=
 github.com/oklog/ulid v1.3.1 h1:EGfNDEx6MqHz8B3uNV6QAib1UR2Lm97sHi3ocA6ESJ4=
@@ -154,11 +142,10 @@ go.uber.org/zap v1.18.1 h1:CSUJ2mjFszzEWt4CdKISEuChVIXGBn3lAPwkRGyVrc4=
 go.uber.org/zap v1.18.1/go.mod h1:xg/QME4nWcxGxrpdeYfq7UvYrLh66cuVKdrbD1XF/NI=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/exp v0.0.0-20220827204233-334a2380cb91 h1:tnebWN09GYg9OLPss1KXj8txwZc6X6uMr6VFdcGNbHw=
-golang.org/x/exp v0.0.0-20220827204233-334a2380cb91/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=
+golang.org/x/exp v0.0.0-20230206171751-46f607a40771 h1:xP7rWLUr1e1n2xkK5YB4LI0hPEy3LJC6Wk+D4pGlOJg=
+golang.org/x/exp v0.0.0-20230206171751-46f607a40771/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc=
 golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs=
 golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
-golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.8.0 h1:LUYupSeNrTNCGzR/hVBk2NHZO4hXcVaW1k4Qx7rjPx8=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
@@ -198,8 +185,8 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f h1:uF6paiQQebLeSXkrTqHqz0MXhXXS1KgF41eUdBNvxK0=
-golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f/go.mod h1:K8+ghG5WaK9qNqU5K3HdILfMLy1f3aNYFI/wnl100a8=
+golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 h1:H2TDz8ibqkAF6YGhCdN3jS9O0/s90v0rJh3X/OLHEUk=
+golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2/go.mod h1:K8+ghG5WaK9qNqU5K3HdILfMLy1f3aNYFI/wnl100a8=
 gonum.org/v1/gonum v0.12.0 h1:xKuo6hzt+gMav00meVPUlXwSdoEJP46BR+wdxQEFK2o=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
diff --git a/query/memory.go b/query/memory.go
index 1d7af86..6fe129c 100644
--- a/query/memory.go
+++ b/query/memory.go
@@ -1,6 +1,9 @@
 package query
 
 import (
+	"fmt"
+	"runtime/debug"
+	"sync"
 	"sync/atomic"
 
 	"github.com/apache/arrow/go/v13/arrow/memory"
@@ -26,12 +29,16 @@ func NewLimitAllocator(limit int64, allocator memory.Allocator) *LimitAllocator
 }
 
 func (a *LimitAllocator) Allocate(size int) []byte {
-	allocated := a.allocated.Add(int64(size))
-	if allocated > a.limit {
-		panic(PanicMemoryLimit)
-	}
+	for {
+		allocated := a.allocated.Load()
+		if allocated+int64(size) > a.limit {
+			panic(PanicMemoryLimit)
+		}
 
-	return a.allocator.Allocate(size)
+		if a.allocated.CompareAndSwap(allocated, allocated+int64(size)) {
+			return a.allocator.Allocate(size)
+		}
+	}
 }
 
 func (a *LimitAllocator) Reallocate(size int, b []byte) []byte {
@@ -39,12 +46,18 @@ func (a *LimitAllocator) Reallocate(size int, b []byte) []byte {
 		return b
 	}
 
-	allocated := a.allocated.Add(int64(size - len(b)))
-	if allocated > a.limit {
-		panic(PanicMemoryLimit)
+	diff := int64(size - len(b))
+	for {
+		allocated := a.allocated.Load()
+		if allocated+diff > a.limit {
+			debug.PrintStack()
+			panic(PanicMemoryLimit)
+		}
+
+		if a.allocated.CompareAndSwap(allocated, allocated+diff) {
+			return a.allocator.Reallocate(size, b)
+		}
 	}
-	buf := a.allocator.Reallocate(size, b)
-	return buf
 }
 
 func (a *LimitAllocator) Free(b []byte) {
@@ -55,3 +68,73 @@ func (a *LimitAllocator) Free(b []byte) {
 func (a *LimitAllocator) Allocated() int {
 	return int(a.allocated.Load())
 }
+
+type allocation struct {
+	stack string
+	size  int
+}
+
+type DebugAllocator struct {
+	sync.Mutex
+
+	allocations map[*byte]allocation
+	allocator   memory.Allocator
+}
+
+func NewDebugAllocator(allocator memory.Allocator) *DebugAllocator {
+	return &DebugAllocator{
+		allocations: make(map[*byte]allocation),
+		allocator:   allocator,
+	}
+}
+
+func (a *DebugAllocator) Allocate(size int) []byte {
+	a.Lock()
+	defer a.Unlock()
+
+	b := a.allocator.Allocate(size)
+	a.allocations[&b[0]] = allocation{
+		stack: string(debug.Stack()),
+		size:  size,
+	}
+	return b
+}
+
+func (a *DebugAllocator) Reallocate(size int, b []byte) []byte {
+	a.Lock()
+	defer a.Unlock()
+
+	delete(a.allocations, &b[0])
+	b = a.allocator.Reallocate(size, b)
+	a.allocations[&b[0]] = allocation{
+		stack: string(debug.Stack()),
+		size:  size,
+	}
+	return b
+}
+
+func (a *DebugAllocator) Free(b []byte) {
+	a.Lock()
+	defer a.Unlock()
+
+	if len(b) == 0 {
+		return
+	}
+
+	delete(a.allocations, &b[0])
+	a.allocator.Free(b)
+}
+
+func (a *DebugAllocator) String() string {
+	a.Lock()
+	defer a.Unlock()
+
+	s := fmt.Sprintf("Allocations remaining: %v\n", len(a.allocations))
+	for _, stack := range a.allocations {
+		s += fmt.Sprintf("Size: %v\n", stack.size)
+		s += stack.stack
+		s += "\n================================================================\n"
+	}
+
+	return s
+}
diff --git a/query/physicalplan/aggregate.go b/query/physicalplan/aggregate.go
index 51764b6..b054743 100644
--- a/query/physicalplan/aggregate.go
+++ b/query/physicalplan/aggregate.go
@@ -200,6 +200,23 @@ func NewHashAggregate(
 	}
 }
 
+func (a *HashAggregate) Close() {
+	for _, arr := range a.groupByArrays {
+		arr.Release()
+	}
+	for _, aggregate := range a.aggregates {
+		for _, aggregation := range aggregate.aggregations {
+			for _, bldr := range aggregation.arrays {
+				bldr.Release()
+			}
+		}
+		for _, bldr := range aggregate.groupByCols {
+			bldr.Release()
+		}
+	}
+	a.next.Close()
+}
+
 func (a *HashAggregate) SetNext(next PhysicalPlan) {
 	a.next = next
 }
@@ -556,22 +573,6 @@ func (a *HashAggregate) Finish(ctx context.Context) error {
 	span.SetAttributes(attribute.Bool("finalStage", a.finalStage))
 	defer span.End()
 
-	defer func() {
-		for _, arr := range a.groupByArrays {
-			arr.Release()
-		}
-		for _, aggregate := range a.aggregates {
-			for _, aggregation := range aggregate.aggregations {
-				for _, bldr := range aggregation.arrays {
-					bldr.Release()
-				}
-			}
-			for _, bldr := range aggregate.groupByCols {
-				bldr.Release()
-			}
-		}
-	}()
-
 	totalRows := 0
 	for _, aggregate := range a.aggregates {
 		if err := a.finishAggregate(ctx, aggregate); err != nil {
diff --git a/query/physicalplan/distinct.go b/query/physicalplan/distinct.go
index 5e50190..738e592 100644
--- a/query/physicalplan/distinct.go
+++ b/query/physicalplan/distinct.go
@@ -62,6 +62,10 @@ func (d *Distinction) Finish(ctx context.Context) error {
 	return d.next.Finish(ctx)
 }
 
+func (d *Distinction) Close() {
+	d.next.Close()
+}
+
 func (d *Distinction) Callback(ctx context.Context, r arrow.Record) error {
 	// Generates high volume of spans. Comment out if needed during development.
 	// ctx, span := d.tracer.Start(ctx, "Distinction/Callback")
@@ -159,7 +163,6 @@ func (d *Distinction) Callback(ctx context.Context, r arrow.Record) error {
 		rows,
 	)
 
-	err := d.next.Callback(ctx, distinctRecord)
-	distinctRecord.Release()
-	return err
+	defer distinctRecord.Release()
+	return d.next.Callback(ctx, distinctRecord)
 }
diff --git a/query/physicalplan/filter.go b/query/physicalplan/filter.go
index 5475ac1..59de5b4 100644
--- a/query/physicalplan/filter.go
+++ b/query/physicalplan/filter.go
@@ -231,6 +231,10 @@ func (f *PredicateFilter) SetNext(next PhysicalPlan) {
 	f.next = next
 }
 
+func (f *PredicateFilter) Close() {
+	f.next.Close()
+}
+
 func (f *PredicateFilter) Callback(ctx context.Context, r arrow.Record) error {
 	// Generates high volume of spans. Comment out if needed during development.
 	// ctx, span := f.tracer.Start(ctx, "PredicateFilter/Callback")
diff --git a/query/physicalplan/ordered_aggregate.go b/query/physicalplan/ordered_aggregate.go
index 80c1b39..b03e0e5 100644
--- a/query/physicalplan/ordered_aggregate.go
+++ b/query/physicalplan/ordered_aggregate.go
@@ -133,6 +133,10 @@ func NewOrderedAggregate(
 	return o
 }
 
+func (a *OrderedAggregate) Close() {
+	a.next.Close()
+}
+
 func (a *OrderedAggregate) SetNext(next PhysicalPlan) {
 	a.next = next
 }
diff --git a/query/physicalplan/ordered_synchronizer.go b/query/physicalplan/ordered_synchronizer.go
index 440c89d..050ee97 100644
--- a/query/physicalplan/ordered_synchronizer.go
+++ b/query/physicalplan/ordered_synchronizer.go
@@ -52,6 +52,10 @@ func NewOrderedSynchronizer(pool memory.Allocator, inputs int, orderByExprs []lo
 	return o
 }
 
+func (o *OrderedSynchronizer) Close() {
+	o.next.Close()
+}
+
 func (o *OrderedSynchronizer) Callback(ctx context.Context, r arrow.Record) error {
 	o.sync.mtx.Lock()
 	o.sync.data = append(o.sync.data, r)
diff --git a/query/physicalplan/physicalplan.go b/query/physicalplan/physicalplan.go
index a3bcf04..805ddb3 100644
--- a/query/physicalplan/physicalplan.go
+++ b/query/physicalplan/physicalplan.go
@@ -24,6 +24,7 @@ type PhysicalPlan interface {
 	Finish(ctx context.Context) error
 	SetNext(next PhysicalPlan)
 	Draw() *Diagram
+	Close()
 }
 
 type ScanPhysicalPlan interface {
@@ -77,6 +78,8 @@ func (e *OutputPlan) Finish(_ context.Context) error {
 	return nil
 }
 
+func (e *OutputPlan) Close() {}
+
 func (e *OutputPlan) SetNext(_ PhysicalPlan) {
 	// OutputPlan should be the last step.
 	// If this gets called we're doing something wrong.
@@ -119,6 +122,11 @@ func (s *TableScan) Execute(ctx context.Context, pool memory.Allocator) error {
 	for _, plan := range s.plans {
 		callbacks = append(callbacks, plan.Callback)
 	}
+	defer func() { // Close all plans to ensure memory cleanup.
+		for _, plan := range s.plans {
+			plan.Close()
+		}
+	}()
 
 	errg, _ := errgroup.WithContext(ctx)
 	errg.Go(recovery.Do(func() error {
@@ -147,7 +155,11 @@ func (s *TableScan) Execute(ctx context.Context, pool memory.Allocator) error {
 		}))
 	}
 
-	return errg.Wait()
+	if err := errg.Wait(); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 type SchemaScan struct {
@@ -209,6 +221,10 @@ type noopOperator struct {
 	next PhysicalPlan
 }
 
+func (p *noopOperator) Close() {
+	p.next.Close()
+}
+
 func (p *noopOperator) Callback(ctx context.Context, r arrow.Record) error {
 	return p.next.Callback(ctx, r)
 }
diff --git a/query/physicalplan/project.go b/query/physicalplan/project.go
index a64da76..bbffabf 100644
--- a/query/physicalplan/project.go
+++ b/query/physicalplan/project.go
@@ -223,6 +223,10 @@ func Project(mem memory.Allocator, tracer trace.Tracer, exprs []logicalplan.Expr
 	return p, nil
 }
 
+func (p *Projection) Close() {
+	p.next.Close()
+}
+
 func (p *Projection) Callback(ctx context.Context, r arrow.Record) error {
 	// Generates high volume of spans. Comment out if needed during development.
 	// ctx, span := p.tracer.Start(ctx, "Projection/Callback")
diff --git a/query/physicalplan/synchronize.go b/query/physicalplan/synchronize.go
index c771be0..7298ac1 100644
--- a/query/physicalplan/synchronize.go
+++ b/query/physicalplan/synchronize.go
@@ -60,3 +60,7 @@ func (m *Synchronizer) SetNextPlan(nextPlan PhysicalPlan) {
 func (m *Synchronizer) Draw() *Diagram {
 	return &Diagram{Details: "Synchronizer", Child: m.next.Draw()}
 }
+
+func (m *Synchronizer) Close() {
+	m.next.Close()
+}
diff --git a/table.go b/table.go
index 71c3f99..944e4ba 100644
--- a/table.go
+++ b/table.go
@@ -914,17 +914,17 @@ func (t *Table) Iterator(
 						if r == nil || r.NumRows() == 0 {
 							return nil
 						}
+						defer r.Release()
 						if err := callback(ctx, r); err != nil {
 							return err
 						}
-						r.Release()
 						return nil
 					}
 
 					switch t := rg.(type) {
 					case arrow.Record:
+						defer t.Release()
 						err := callback(ctx, t)
-						t.Release()
 						if err != nil {
 							return err
 						}
@@ -937,11 +937,14 @@ func (t *Table) Iterator(
 							continue
 						}
 						if converter.NumRows() >= bufferSize {
-							r := converter.NewRecord()
-							if err := callback(ctx, r); err != nil {
+							err := func() error {
+								r := converter.NewRecord()
+								defer r.Release()
+								return callback(ctx, r)
+							}()
+							if err != nil {
 								return err
 							}
-							r.Release()
 						}
 					default:
 						return fmt.Errorf("unknown row group type: %T", t)


Test patch:
diff --git a/db_test.go b/db_test.go
index eccb60a..df8acd4 100644
--- a/db_test.go
+++ b/db_test.go
@@ -1776,24 +1776,30 @@ func Test_DB_Limiter(t *testing.T) {
 	_, err = table.InsertBuffer(ctx, buf)
 	require.NoError(t, err)
 
-	pool := query.NewLimitAllocator(1024*1024*1024, memory.DefaultAllocator)
-	engine := query.NewEngine(pool, db.TableProvider())
-	err = engine.ScanTable("test").
-		Filter(
-			logicalplan.And(
-				logicalplan.Col("labels.namespace").Eq(logicalplan.Literal("default")),
-			),
-		).
-		Aggregate(
-			[]logicalplan.Expr{logicalplan.Sum(logicalplan.Col("value"))},
-			[]logicalplan.Expr{logicalplan.Col("labels.namespace")},
-		).
-		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
-			require.Equal(t, int64(1), r.NumRows())
-			return nil
-		})
-	require.NoError(t, err)
+	for i := 0; i < 1024; i++ {
+		t.Run(fmt.Sprintf("limit-%v", i), func(t *testing.T) {
+			debug := query.NewDebugAllocator(memory.DefaultAllocator)
+			pool := query.NewLimitAllocator(int64(i), debug)
+			engine := query.NewEngine(pool, db.TableProvider())
+			err = engine.ScanTable("test").
+				Filter(
+					logicalplan.And(
+						logicalplan.Col("labels.namespace").Eq(logicalplan.Literal("default")),
+					),
+				).
+				Aggregate(
+					[]logicalplan.Expr{logicalplan.Sum(logicalplan.Col("value"))},
+					[]logicalplan.Expr{logicalplan.Col("labels.namespace")},
+				).
+				Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+					return nil
+				})
 
-	// Expect no bytes allocated after the query has finished
-	require.Equal(t, 0, pool.Allocated())
+			// Expect no bytes allocated after the query has finished
+			if pool.Allocated() != 0 {
+				fmt.Println(debug)
+			}
+			require.Equal(t, 0, pool.Allocated())
+		})
+	}
 }
diff --git a/query/physicalplan/physicalplan_test.go b/query/physicalplan/physicalplan_test.go
index 298a91c..fdfb5d5 100644
--- a/query/physicalplan/physicalplan_test.go
+++ b/query/physicalplan/physicalplan_test.go
@@ -104,3 +104,7 @@ func (m *mockPhysicalPlan) SetNext(next PhysicalPlan) {
 func (m *mockPhysicalPlan) Draw() *Diagram {
 	return &Diagram{}
 }
+
+func (m *mockPhysicalPlan) Close() {
+	m.next.Close()
+}