From 623aaec9d896cc8dbe9421bfd0a3c285107a9b4e Mon Sep 17 00:00:00 2001
From: Ben Visness <bvisness@gmail.com>
Date: Thu, 21 Oct 2021 01:42:34 -0500
Subject: [PATCH] Ensure that that one goroutine exits when the iterator is
 closed

This resolves that completely nonsensical memory leak situation. As far
as we can understand, the cause was a hodgepodge of the following:

- There is some buffer sharing going on deep in pgx
- Queries made with a cancellable but long-running context (like that
used for background jobs) would leave iterator-related goroutines
hanging
- These goroutines had a pgx `rows` object in their closures, preventing
the row stuff from being garbage collected
- If you look at a profile, it all appears to be caused by whatever
functions were doing the most database queries / reading the most from
Postgres. In fact those things were _allocating_ the most but not
retaining any of that data - it was being retained by these other
goroutines because of magic buffer sharing huzzah I love it

We could have solved this in approximately 30 minutes if Go could
actually tell us what is keeping things alive in the heap, instead of
just tracking allocations.
---
 src/db/db.go | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/db/db.go b/src/db/db.go
index 69defd5..32996e2 100644
--- a/src/db/db.go
+++ b/src/db/db.go
@@ -93,11 +93,13 @@ type StructQueryIterator struct {
 	fieldPaths [][]int
 	rows       pgx.Rows
 	destType   reflect.Type
+	closed     chan struct{}
 }
 
 func (it *StructQueryIterator) Next() (interface{}, bool) {
 	hasNext := it.rows.Next()
 	if !hasNext {
+		it.Close()
 		return nil, false
 	}
 
@@ -169,6 +171,10 @@ func (it *StructQueryIterator) Next() (interface{}, bool) {
 
 func (it *StructQueryIterator) Close() {
 	it.rows.Close()
+	select {
+	case it.closed <- struct{}{}:
+	default:
+	}
 }
 
 func (it *StructQueryIterator) ToSlice() []interface{} {
@@ -241,6 +247,7 @@ func Query(ctx context.Context, conn ConnOrTx, destExample interface{}, query st
 		fieldPaths: fieldPaths,
 		rows:       rows,
 		destType:   destType,
+		closed:     make(chan struct{}, 1),
 	}
 
 	// Ensure that iterators are closed if context is cancelled. Otherwise, iterators can hold
@@ -250,8 +257,11 @@ func Query(ctx context.Context, conn ConnOrTx, destExample interface{}, query st
 		if done == nil {
 			return
 		}
-		<-done
-		it.Close()
+		select {
+		case <-done:
+			it.Close()
+		case <-it.closed:
+		}
 	}()
 
 	return it, nil