store: fix nil panic in proxy heap (thanos-io#5746)

* store: fix nil panic in proxy heap

With lazy proxying we need to wait for at least one response before we
can build a heap properly.

Closes thanos-io#5717
Fixes thanos-io#5552.

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>

* store: panic if At() called without Next()

This is pretty much a bug if this ever happens so complain loudly.

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>

* store: adjust test after recent changes

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>
Signed-off-by: utukj <utukphd@gmail.com>

pkg/store/proxy_heap.go

@@ -112,7 +112,7 @@ func (d *dedupResponseHeap) Next() bool {
 
 func (d *dedupResponseHeap) At() *storepb.SeriesResponse {
 	if len(d.responses) == 0 {
-		return nil
+		panic("BUG: At() called with no responses; please call At() only if Next() returns true")
 	} else if len(d.responses) == 1 {
 		return d.responses[0]
 	}
@@ -293,6 +293,19 @@ func (l *lazyRespSet) Empty() bool {
 	l.bufferedResponsesMtx.Lock()
 	defer l.bufferedResponsesMtx.Unlock()
 
+	// NOTE(GiedriusS): need to wait here for at least one
+	// response so that we could build the heap properly.
+	if l.noMoreData && len(l.bufferedResponses) == 0 {
+		return true
+	}
+
+	for len(l.bufferedResponses) == 0 {
+		l.dataOrFinishEvent.Wait()
+		if l.noMoreData && len(l.bufferedResponses) == 0 {
+			break
+		}
+	}
+
 	return len(l.bufferedResponses) == 0 && l.noMoreData
 }
 

pkg/store/proxy_test.go

@@ -1991,7 +1991,15 @@ func TestDedupRespHeap_Deduplication(t *testing.T) {
 			responses: []*storepb.SeriesResponse{},
 			testFn: func(responses []*storepb.SeriesResponse, h *dedupResponseHeap) {
 				testutil.Equals(t, false, h.Next())
-				testutil.Equals(t, (*storepb.SeriesResponse)(nil), h.At())
+
+				callAtExpectPanic := func() {
+					defer func() {
+						testutil.Assert(t, recover() != nil, "expected a panic from At()")
+					}()
+
+					h.At()
+				}
+				callAtExpectPanic()
 			},
 		},
 		{

test/e2e/query_test.go

@@ -1560,4 +1560,9 @@ func TestConnectedQueriesWithLazyProxy(t *testing.T) {
 		return "sum(up)"
 	}, time.Now, promclient.QueryOptions{}, 1)
 	testutil.Equals(t, model.SampleValue(1.0), result[0].Value)
+
+	instantQuery(t, context.Background(), querier2.Endpoint("http"), func() string {
+		return "sum(metric_that_does_not_exist)"
+	}, time.Now, promclient.QueryOptions{}, 0)
+
 }