Repository: segmentio/parquet-go
Pull request #337 (2022-09-06 14:12:17)

close rowGroupCursors when Popped due to io.EOF in mergedRowGroupRows
This commit fixes a potential goroutine leak. mergeRowGroupRows cursors, when
initialized, might spawn a goroutine (e.g. asyncPages.readPages). Previously,
when a cursor returned io.EOF, this cursor was popped from the min-heap but
never closed, resulting in leaked goroutines.

cc @achille-roussel @Pryz 

asubiotto: Any pointers on where to add a regression test?
joe-elliott: We are possibly seeing this in a rollout today. We have ~100s of orphaned goroutines parked here:

https://github.com/segmentio/parquet-go/blob/main/page.go#L222

The goroutines do seem to be slowly falling off. Currently still investigating.

![image](https://user-images.githubusercontent.com/2272392/189987170-c45ccc3f-1897-4b73-a4da-f13a7c9c0a69.png)

Is that consistent with this issue?



achille-roussel: If you are using the Merge* functions it would be related. 

Alternatively you might be missing a Close call on Rows instances in the code using parquet-go.
joe-elliott: Yup, this was my bad. Saw some new behavior on a rollout and was concerned, but this is attributable to other changes in the rollout. I'll leave my original post only for continuity of the conversation :P
asubiotto: Updated and added test.
asubiotto: @joe-elliott there are some finalizers set up to close resources on GC (including the goroutines leaked by the behavior this PR fixes). I found it helpful to disable those and run a leak detector to track down this leak.

Patch:
diff --git a/merge.go b/merge.go
index aef1410..6e6b8dd 100644
--- a/merge.go
+++ b/merge.go
@@ -120,7 +120,12 @@ func (r *mergedRowGroupRows) ReadRows(rows []Row) (n int, err error) {
 				r.err = err
 				return n, err
 			}
+			c := r.cursors[0]
 			heap.Pop(r)
+			if err := c.close(); err != nil {
+				r.err = err
+				return n, err
+			}
 		} else {
 			heap.Fix(r, 0)
 		}


Test patch:
diff --git a/merge_test.go b/merge_test.go
index 4d877d5..5be47a6 100644
--- a/merge_test.go
+++ b/merge_test.go
@@ -17,6 +17,75 @@ const (
 	rowsPerGroup = benchmarkNumRows
 )
 
+type wrappedRowGroup struct {
+	parquet.RowGroup
+	rowsCallback func(parquet.Rows) parquet.Rows
+}
+
+func (r wrappedRowGroup) Rows() parquet.Rows {
+	return r.rowsCallback(r.RowGroup.Rows())
+}
+
+type wrappedRows struct {
+	parquet.Rows
+	closed bool
+}
+
+func (r *wrappedRows) Close() error {
+	r.closed = true
+	return r.Rows.Close()
+}
+
+func TestMergeRowGroupsCursorsAreClosed(t *testing.T) {
+
+	type model struct {
+		A int
+	}
+
+	schema := parquet.SchemaOf(model{})
+	options := []parquet.RowGroupOption{
+		parquet.SortingColumns(
+			parquet.Ascending(schema.Columns()[0]...),
+		),
+	}
+
+	prng := rand.New(rand.NewSource(0))
+	rowGroups := make([]parquet.RowGroup, numRowGroups)
+	rows := make([]*wrappedRows, 0, numRowGroups)
+
+	for i := range rowGroups {
+		rowGroups[i] = wrappedRowGroup{
+			RowGroup: sortedRowGroup(options, randomRowsOf(prng, rowsPerGroup, model{})...),
+			rowsCallback: func(r parquet.Rows) parquet.Rows {
+				wrapped := &wrappedRows{Rows: r}
+				rows = append(rows, wrapped)
+				return wrapped
+			},
+		}
+	}
+
+	m, err := parquet.MergeRowGroups(rowGroups, options...)
+	if err != nil {
+		t.Fatal(err)
+	}
+	func() {
+		mergedRows := m.Rows()
+		defer mergedRows.Close()
+
+		// Add 1 more slot to the buffer to force an io.EOF on the first read.
+		rbuf := make([]parquet.Row, (numRowGroups*rowsPerGroup)+1)
+		if _, err := mergedRows.ReadRows(rbuf); !errors.Is(err, io.EOF) {
+			t.Fatal(err)
+		}
+	}()
+
+	for i, wrapped := range rows {
+		if !wrapped.closed {
+			t.Fatalf("RowGroup %d not closed", i)
+		}
+	}
+}
+
 func BenchmarkMergeRowGroups(b *testing.B) {
 	for _, test := range readerTests {
 		b.Run(test.scenario, func(b *testing.B) {