From b954c00fa8136c87870302ad0297e704f280e2c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Nordstr=C3=B6m?= <erik@timescale.com>
Date: Fri, 11 Mar 2022 17:46:05 +0100
Subject: [PATCH] Fix memory handling during scans

Scan functions cannot be called on a per-tuple memory context as they
might allocate data that need to live until the end of the scan. Fix
this in a couple of places to ensure correct memory handling.

Fixes #4148, #4145
---
 CHANGELOG.md          |  4 ++++
 src/chunk_scan.c      |  2 +-
 src/scan_iterator.c   |  2 ++
 src/telemetry/stats.c | 42 +++++++++++++++++++-----------------------
 4 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 78a4b8fe2..a693015e9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,10 @@ accidentally triggering the load of a previous DB version.**
 
 **Bugfixes**
 * #4122 Fix segfault on INSERT into distributed hypertable
+* #4161 Fix memory handling during scans
+
+**Thanks**
+* @abrownsword for reporting a crash in the telemetry reporter
 
 ## 2.6.0 (2022-02-16)
 This release is medium priority for upgrade. We recommend that you upgrade at the next available opportunity.
diff --git a/src/chunk_scan.c b/src/chunk_scan.c
index 5cdb7120e..4f4b6c489 100644
--- a/src/chunk_scan.c
+++ b/src/chunk_scan.c
@@ -252,9 +252,9 @@ ts_chunk_scan_by_constraints(const Hyperspace *hs, const List *dimension_vecs,
 				prev_chunk_oid = chunk->table_id;
 			}
 
+			MemoryContextSwitchTo(work_mcxt);
 			/* Only one chunk should match */
 			Assert(ts_scan_iterator_next(&chunk_it) == NULL);
-			MemoryContextSwitchTo(work_mcxt);
 		}
 	}
 
diff --git a/src/scan_iterator.c b/src/scan_iterator.c
index 1742f3d50..f2a1d6e06 100644
--- a/src/scan_iterator.c
+++ b/src/scan_iterator.c
@@ -56,5 +56,7 @@ ts_scan_iterator_scan_key_init(ScanIterator *iterator, AttrNumber attributeNumbe
 TSDLLEXPORT void
 ts_scan_iterator_rescan(ScanIterator *iterator)
 {
+	MemoryContext oldmcxt = MemoryContextSwitchTo(iterator->scankey_mcxt);
 	ts_scanner_rescan(&iterator->ctx, NULL);
+	MemoryContextSwitchTo(oldmcxt);
 }
diff --git a/src/telemetry/stats.c b/src/telemetry/stats.c
index ce23e111f..35084e166 100644
--- a/src/telemetry/stats.c
+++ b/src/telemetry/stats.c
@@ -27,7 +27,6 @@ typedef struct StatsContext
 {
 	TelemetryStats *stats;
 	ScanIterator compressed_chunk_stats_iterator;
-	bool iterator_valid;
 } StatsContext;
 
 /*
@@ -338,27 +337,25 @@ static bool
 get_chunk_compression_stats(StatsContext *statsctx, const Chunk *chunk,
 							Form_compression_chunk_size compr_stats)
 {
+	TupleInfo *ti;
+	MemoryContext oldmcxt;
+
 	if (!ts_chunk_is_compressed(chunk))
 		return false;
 
+	/* Need to execute the scan functions on the long-lived memory context */
+	oldmcxt = MemoryContextSwitchTo(statsctx->compressed_chunk_stats_iterator.scankey_mcxt);
 	ts_scan_iterator_scan_key_reset(&statsctx->compressed_chunk_stats_iterator);
 	ts_scan_iterator_scan_key_init(&statsctx->compressed_chunk_stats_iterator,
 								   Anum_compression_chunk_size_pkey_chunk_id,
 								   BTEqualStrategyNumber,
 								   F_INT4EQ,
 								   Int32GetDatum(chunk->fd.id));
+	ts_scan_iterator_start_or_restart_scan(&statsctx->compressed_chunk_stats_iterator);
+	ti = ts_scan_iterator_next(&statsctx->compressed_chunk_stats_iterator);
+	MemoryContextSwitchTo(oldmcxt);
 
-	if (statsctx->iterator_valid)
-	{
-		ts_scan_iterator_rescan(&statsctx->compressed_chunk_stats_iterator);
-	}
-	else
-	{
-		ts_scan_iterator_start_scan(&statsctx->compressed_chunk_stats_iterator);
-		statsctx->iterator_valid = true;
-	}
-
-	if (ts_scan_iterator_next(&statsctx->compressed_chunk_stats_iterator))
+	if (ti)
 	{
 		Form_compression_chunk_size fd;
 		bool should_free;
@@ -379,9 +376,8 @@ get_chunk_compression_stats(StatsContext *statsctx, const Chunk *chunk,
 	/*
 	 * Should only get here if a compressed chunk is missing stats for some
 	 * reason. The iterator will automatically close if no tuple is found, so
-	 * need to make sure it is re-opened next time this function is called.
+	 * it will be re-opened next time this function is called.
 	 */
-	statsctx->iterator_valid = false;
 
 	return false;
 }
@@ -496,7 +492,6 @@ ts_telemetry_stats_gather(TelemetryStats *stats)
 	MemoryContext oldmcxt, relmcxt;
 	StatsContext statsctx = {
 		.stats = stats,
-		.iterator_valid = false,
 	};
 
 	MemSet(stats, 0, sizeof(*stats));
@@ -511,12 +506,6 @@ ts_telemetry_stats_gather(TelemetryStats *stats)
 
 	relmcxt = AllocSetContextCreate(CurrentMemoryContext, "RelationStats", ALLOCSET_DEFAULT_SIZES);
 
-	/*
-	 * Use temporary per-tuple memory context to not accumulate cruft during
-	 * processing of pg_class.
-	 */
-	oldmcxt = MemoryContextSwitchTo(relmcxt);
-
 	while (true)
 	{
 		HeapTuple tup;
@@ -526,7 +515,6 @@ ts_telemetry_stats_gather(TelemetryStats *stats)
 		const Hypertable *ht = NULL;
 		const ContinuousAgg *cagg = NULL;
 
-		MemoryContextReset(relmcxt);
 		tup = systable_getnext(scan);
 
 		if (!HeapTupleIsValid(tup))
@@ -537,6 +525,13 @@ ts_telemetry_stats_gather(TelemetryStats *stats)
 		if (should_ignore_relation(catalog, class))
 			continue;
 
+		/*
+		 * Use temporary per-relation memory context to not accumulate cruft
+		 * during processing of pg_class.
+		 */
+		oldmcxt = MemoryContextSwitchTo(relmcxt);
+		MemoryContextReset(relmcxt);
+
 		reltype = classify_relation(class, htcache, &ht, &chunk, &cagg);
 
 		switch (reltype)
@@ -592,9 +587,10 @@ ts_telemetry_stats_gather(TelemetryStats *stats)
 			case RELTYPE_OTHER:
 				break;
 		}
+
+		MemoryContextSwitchTo(oldmcxt);
 	}
 
-	MemoryContextSwitchTo(oldmcxt);
 	systable_endscan(scan);
 	table_close(rel, AccessShareLock);
 	ts_scan_iterator_close(&statsctx.compressed_chunk_stats_iterator);