querier: release LSS query result buffers deterministically

pp/entrypoint/bridge/series_data_data_storage.cpp

@@ -1,6 +1,8 @@
 #include "series_data_data_storage.h"
 
 #include <cassert>
+#include <ranges>
+#include <span>
 #include <spanstream>
 
 #include "entrypoint/types/data_storage.h"
@@ -266,6 +268,25 @@ extern "C" void prompp_series_data_data_storage_query_first_timestamps(void* arg
                          [&storage](uint32_t series_id) { return Decoder::get_series_min_timestamp(storage, series_id); });
 }
 
+extern "C" void prompp_series_data_data_storage_query_stalenan_series(void* args) {
+  using series_data::Decoder;
+
+  struct Arguments {
+    DataStoragePtr data_storage;
+    SliceView<LabelSetID> series_ids;
+    entrypoint::types::StaleNaNSeriesWithGoLabels* series;
+  };
+
+  const auto in = static_cast<Arguments*>(args);
+  const auto& storage = *in->data_storage;
+
+  auto series = std::span(in->series, in->series_ids.size());
+  for (auto&& [stalenan_series, series_id] : std::ranges::views::zip(series, in->series_ids)) {
+    stalenan_series.timestamp = Decoder::get_series_min_timestamp(storage, series_id);
+    stalenan_series.series_id = series_id;
+  }
+}
+
 extern "C" void prompp_series_data_data_storage_allocated_memory(void* args, void* res) {
   struct Arguments {
     DataStoragePtr data_storage;

pp/entrypoint/bridge/series_data_data_storage.h

@@ -181,6 +181,18 @@ void prompp_series_data_data_storage_instant_query(void* args, void* res);
  */
 void prompp_series_data_data_storage_query_first_timestamps(void* args, void* res);
 
+/**
+ * @brief Fill stalenan series (first sample timestamp + series id) per series id.
+ *
+ * @param args {
+ *        dataStorage uintptr  // pointer to constructed data storage
+ *        seriesIds   []uint32 // series ids
+ *        series      uintptr  // pointer to []querier.StaleNaNSeries (same length as seriesIds);
+ *                             // timestamp and seriesID fields are filled from storage
+ * }
+ */
+void prompp_series_data_data_storage_query_stalenan_series(void* args);
+
 /**
  * @brief finishes all Queriers after data load.
  *

pp/entrypoint/types/querier.h

@@ -58,6 +58,20 @@ class InstantQuerierWithArgumentsWrapper {
 };
 
 struct SampleWithGoLabels : public ::series_data::encoder::Sample {
+  // series_id is filled by the instant querier alongside the sample; go_labels_ is written
+  // from the Go side (querier.InstantSeries.LabelSet). This layout mirrors the Go struct.
+  uint32_t series_id;
+
+ private:
+  char go_labels_[Sizeof_GoLabels];
+};
+
+struct StaleNaNSeriesWithGoLabels {
+  // timestamp and series_id are filled by the data storage stalenan query; go_labels_ is
+  // written from the Go side (querier.StaleNaNSeries.labelSet). This layout mirrors the Go struct.
+  PromPP::Primitives::Timestamp timestamp;
+  uint32_t series_id;
+
  private:
   char go_labels_[Sizeof_GoLabels];
 };

pp/go/cppbridge/data_storage.go

@@ -126,6 +126,13 @@ func (ds *DataStorage) QueryFirstTimestamps(seriesIDs []uint32, timestamps []int
 	runtime.KeepAlive(ds)
 }
 
+// QueryStaleNaNSeries fills the first sample timestamp (Prometheus ms) and the series id for each
+// series in seriesIDs directly into the C-shared series slice (pointed to by series).
+func (ds *DataStorage) QueryStaleNaNSeries(seriesIDs []uint32, series uintptr) {
+	seriesDataDataStorageQueryStaleNaNSeries(ds.dataStorage, seriesIDs, series)
+	runtime.KeepAlive(ds)
+}
+
 func (ds *DataStorage) QueryFinal(queriers []uintptr) {
 	seriesDataDataStorageQueryFinal(queriers)
 	runtime.KeepAlive(queriers)

pp/go/cppbridge/entrypoint.go

@@ -2214,6 +2214,20 @@ func seriesDataDataStorageQueryFirstTimestamps(dataStorage uintptr, seriesIDs []
 	)
 }
 
+func seriesDataDataStorageQueryStaleNaNSeries(dataStorage uintptr, seriesIDs []uint32, series uintptr) {
+	args := struct {
+		dataStorage uintptr
+		seriesIDs   []uint32
+		series      uintptr
+	}{dataStorage, seriesIDs, series}
+
+	testGC()
+	fastcgo.UnsafeCall1(
+		C.prompp_series_data_data_storage_query_stalenan_series,
+		uintptr(unsafe.Pointer(&args)),
+	)
+}
+
 func seriesDataDataStorageQueryFinal(queriers []uintptr) {
 	args := struct {
 		queriers []uintptr

pp/go/cppbridge/entrypoint.h

@@ -1746,6 +1746,18 @@ void prompp_series_data_data_storage_instant_query(void* args, void* res);
  */
 void prompp_series_data_data_storage_query_first_timestamps(void* args, void* res);
 
+/**
+ * @brief Fill stalenan series (first sample timestamp + series id) per series id.
+ *
+ * @param args {
+ *        dataStorage uintptr  // pointer to constructed data storage
+ *        seriesIds   []uint32 // series ids
+ *        series      uintptr  // pointer to []querier.StaleNaNSeries (same length as seriesIds);
+ *                             // timestamp and seriesID fields are filled from storage
+ * }
+ */
+void prompp_series_data_data_storage_query_stalenan_series(void* args);
+
 /**
  * @brief finishes all Queriers after data load.
  *

pp/go/cppbridge/lss_snapshot.go

@@ -158,18 +158,32 @@ func newLSSQueryResult(
 	}
 
 	if status != LSSQueryStatusMatch {
-		primitivesLabelSetMatchesFree(lqr)
+		lqr.Close()
 
 		return lqr
 	}
 
 	runtime.SetFinalizer(lqr, func(result *LSSQueryResult) {
-		primitivesLabelSetMatchesFree(result)
+		result.Close()
 	})
 
 	return lqr
 }
 
+// Close frees the C-allocated result buffers and cancels the finalizer.
+// It is idempotent: subsequent calls (and the finalizer) are no-ops.
+// After Close the result must not be read anymore.
+func (r *LSSQueryResult) Close() {
+	if r.matches == nil && r.labelSetLengths == nil {
+		return
+	}
+
+	runtime.SetFinalizer(r, nil)
+	primitivesLabelSetMatchesFree(r)
+	r.matches = nil
+	r.labelSetLengths = nil
+}
+
 func (r *LSSQueryResult) IndexOf(seriesID uint32) int {
 	for i, match := range r.matches {
 		if match == seriesID {

pp/go/storage/head/shard/data_storage.go

@@ -96,6 +96,14 @@ func (ds *DataStorage) QueryFirstTimestamps(ids []uint32, timestamps []int64) {
 	ds.locker.RUnlock()
 }
 
+// QueryStaleNaNSeries fills the first sample timestamp (Prometheus ms) and the series id for each
+// series in ids directly into the C-shared series slice (pointed to by series).
+func (ds *DataStorage) QueryStaleNaNSeries(ids []uint32, series uintptr) {
+	ds.locker.RLock()
+	ds.dataStorage.QueryStaleNaNSeries(ids, series)
+	ds.locker.RUnlock()
+}
+
 // QueryStatus get head status from [DataStorage].
 func (ds *DataStorage) QueryStatus(status *cppbridge.HeadStatus) {
 	ds.locker.RLock()

pp/go/storage/querier/chunk_querier.go

@@ -142,6 +142,9 @@ func (q *ChunkQuerier[TTask, TDataStorage, TLSS, TShard, THead]) Select(
 	defer poolProvider.PutSnapshots(snapshots)
 	lssQueryResults := poolProvider.GetLSSQueryResults()
 	defer poolProvider.PutLSSQueryResults(lssQueryResults)
+	// The chunk series sets take the series id from the recoded chunks and do not
+	// reference the query result buffers, so they can be released on return.
+	defer closeLSSQueryResults(lssQueryResults)
 
 	if err = queryLss(lssQueryChunkQuerySelector, q.head, matchers, snapshots, lssQueryResults); err != nil {
 		logger.Warnf("[ChunkQuerier]: failed: %s", err)
@@ -162,7 +165,6 @@ func (q *ChunkQuerier[TTask, TDataStorage, TLSS, TShard, THead]) Select(
 		}
 
 		chunkSeriesSets[shardID] = NewChunkSeriesSet(
-			lssQueryResults[shardID],
 			snapshots[shardID],
 			cppbridge.NewSerializedChunkRecoder(serializedData, cppbridge.TimeInterval{MinT: q.mint, MaxT: q.maxt}),
 		)

pp/go/storage/querier/chunk_series.go

@@ -15,25 +15,23 @@ import (
 
 // ChunkSeriesSet contains a set of chunked series.
 type ChunkSeriesSet struct {
-	lssQueryResult   *cppbridge.LSSQueryResult
 	labelSetSnapshot *cppbridge.LabelSetSnapshot
 	chunkRecoder     *cppbridge.ChunkRecoder
 
-	index            int
 	lastRecodedChunk *cppbridge.RecodedChunk
 	chunkSeries      *ChunkSeries
 
 	recoderIsExhausted bool
 }
 
 // NewChunkSeriesSet init new [ChunkSeriesSet].
+// The series id is taken directly from the recoded chunks during iteration, so the
+// query result buffer does not need to be retained (see [cppbridge.LSSQueryResult.Close]).
 func NewChunkSeriesSet(
-	lssQueryResult *cppbridge.LSSQueryResult,
 	labelSetSnapshot *cppbridge.LabelSetSnapshot,
 	chunkRecoder *cppbridge.ChunkRecoder,
 ) *ChunkSeriesSet {
 	return &ChunkSeriesSet{
-		lssQueryResult:   lssQueryResult,
 		labelSetSnapshot: labelSetSnapshot,
 		chunkRecoder:     chunkRecoder,
 	}
@@ -72,24 +70,8 @@ func (css *ChunkSeriesSet) Next() bool {
 		css.lastRecodedChunk = nil
 	}
 
-	var lsID uint32
-
-	for {
-		if css.index >= css.lssQueryResult.Len() {
-			return false
-		}
-
-		lsID, _ = css.lssQueryResult.GetByIndex(css.index)
-
-		if lsID == seriesID {
-			break
-		}
-
-		css.index++
-	}
-
 	css.chunkSeries = &ChunkSeries{
-		labelSet:      labels.NewLabelsWithLSS(css.labelSetSnapshot, lsID),
+		labelSet:      labels.NewLabelsWithLSS(css.labelSetSnapshot, seriesID),
 		recodedChunks: recodedChunks,
 	}
 

pp/go/storage/querier/instant_series.go

@@ -24,22 +24,22 @@ func NewInstantSeriesSlice(size int, defaultTimestamp int64) []InstantSeries {
 
 // InstantSeriesSet contains a instant set of series, allows to iterate over sorted, populated series.
 type InstantSeriesSet struct {
-	lssQueryResult              *cppbridge.LSSQueryResult
 	labelSetSnapshot            *cppbridge.LabelSetSnapshot
 	valueNotFoundTimestampValue int64
 	nextSeriesIndex             int
 	instantSeries               []InstantSeries
 }
 
 // NewInstantSeriesSet init new [InstantSeriesSet].
+// The series id is filled inline in each instantSeries by the C++ instant query (next to the
+// sample), so the query result buffer is not referenced here and can be released right after
+// the data storage query (see [cppbridge.LSSQueryResult.Close]).
 func NewInstantSeriesSet(
-	lssQueryResult *cppbridge.LSSQueryResult,
 	labelSetSnapshot *cppbridge.LabelSetSnapshot,
 	valueNotFoundTimestampValue int64,
 	instantSeries []InstantSeries,
 ) *InstantSeriesSet {
 	return &InstantSeriesSet{
-		lssQueryResult:              lssQueryResult,
 		labelSetSnapshot:            labelSetSnapshot,
 		valueNotFoundTimestampValue: valueNotFoundTimestampValue,
 		instantSeries:               instantSeries,
@@ -70,8 +70,8 @@ func (ss *InstantSeriesSet) Next() bool {
 		ss.nextSeriesIndex++
 	}
 
-	lsID, _ := ss.lssQueryResult.GetByIndex(ss.nextSeriesIndex)
-	ss.instantSeries[ss.nextSeriesIndex].LabelSet = labels.NewLabelsWithLSS(ss.labelSetSnapshot, lsID)
+	series := &ss.instantSeries[ss.nextSeriesIndex]
+	series.LabelSet = labels.NewLabelsWithLSS(ss.labelSetSnapshot, series.SeriesID)
 
 	ss.nextSeriesIndex++
 	return true
@@ -87,9 +87,14 @@ func (*InstantSeriesSet) Warnings() annotations.Annotations {
 //
 
 // InstantSeries is a instant stream of data points belonging to a metric.
+// NOTE: this struct is shared with C++ (entrypoint::types::SampleWithGoLabels):
+// the data storage instant query writes Timestamp/Value into it by pointer, so its
+// layout must stay in sync with the C++ side. Any field added here must be mirrored
+// in SampleWithGoLabels (SeriesID -> series_id_, LabelSet -> go_labels_).
 type InstantSeries struct {
 	Timestamp int64
 	Value     float64
+	SeriesID  uint32
 	LabelSet  labels.Labels
 }
 

pp/go/storage/querier/querier.go

@@ -306,6 +306,9 @@ func (q *Querier[TTask, TDataStorage, TLSS, TShard, THead]) selectInstant(
 	defer poolProvider.PutSnapshots(snapshots)
 	lssQueryResults := poolProvider.GetLSSQueryResults()
 	defer poolProvider.PutLSSQueryResults(lssQueryResults)
+	// The instant series sets copy out the series ids they need, so the query result
+	// buffers can be released on return (after all in-flight C queries have completed).
+	defer closeLSSQueryResults(lssQueryResults)
 
 	if err = queryLss(lssQueryInstantQuerySelector, q.head, matchers, snapshots, lssQueryResults); err != nil {
 		logger.Warnf("[QUERIER]: failed to instant: %s", err)
@@ -342,7 +345,6 @@ func (q *Querier[TTask, TDataStorage, TLSS, TShard, THead]) selectInstant(
 			}
 
 			seriesSets[shardID] = NewInstantSeriesSet(
-				lssQueryResult,
 				snapshots[shardID],
 				valueNotFoundTimestampValue,
 				instantSeries,
@@ -396,6 +398,10 @@ func (q *Querier[TTask, TDataStorage, TLSS, TShard, THead]) selectRange(
 	defer poolProvider.PutSnapshots(snapshots)
 	lssQueryResults := poolProvider.GetLSSQueryResults()
 	defer poolProvider.PutLSSQueryResults(lssQueryResults)
+	// Range/aggr series sets only use the result length at construction, so the query
+	// result buffers can be released on return (queryDataStorage has already waited for
+	// all in-flight C queries by then).
+	defer closeLSSQueryResults(lssQueryResults)
 
 	if err = queryLss(lssQueryRangeQuerySelector, q.head, matchers, snapshots, lssQueryResults); err != nil {
 		logger.Warnf("[QUERIER]: failed to range: %s", err)
@@ -720,6 +726,18 @@ func queryLabelValues[
 	return lvs, anns, nil
 }
 
+// closeLSSQueryResults releases the C-allocated buffers held by the query results.
+// It is safe to call once every series set built from these results has copied out
+// (or no longer needs) the data it references. Nil entries are skipped and Close is
+// idempotent, so the finalizer stays a harmless fallback.
+func closeLSSQueryResults(lssQueryResults []*cppbridge.LSSQueryResult) {
+	for _, lssQueryResult := range lssQueryResults {
+		if lssQueryResult != nil {
+			lssQueryResult.Close()
+		}
+	}
+}
+
 // queryLss returns query results and snapshots.
 //
 //revive:disable-next-line:cyclomatic but readable.