From 99582ba8357a71870ea481a8ef99bde43c375f1c Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Mon, 12 Apr 2021 19:44:03 +0000
Subject: [PATCH] Implement HasReports for live data from collectors

---
 app/multitenant/aws_collector.go |  4 +++
 app/multitenant/collector.go     | 43 ++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/app/multitenant/aws_collector.go b/app/multitenant/aws_collector.go
index 0502e0321..0a387e5c1 100644
--- a/app/multitenant/aws_collector.go
+++ b/app/multitenant/aws_collector.go
@@ -564,6 +564,10 @@ func (c *awsCollector) HasReports(ctx context.Context, timestamp time.Time) (boo
 	if err != nil {
 		return false, err
 	}
+	if time.Since(timestamp) < c.cfg.Window {
+		has, err := c.hasReportsFromLive(ctx, userid)
+		return has, err
+	}
 	start := timestamp.Add(-c.cfg.Window)
 	reportKeys, err := c.getReportKeys(ctx, userid, start, timestamp)
 	return len(reportKeys) > 0, err
diff --git a/app/multitenant/collector.go b/app/multitenant/collector.go
index b6bcc7107..25693d711 100644
--- a/app/multitenant/collector.go
+++ b/app/multitenant/collector.go
@@ -3,6 +3,7 @@ package multitenant
 // Collect reports from probes per-tenant, and supply them to queriers on demand
 
 import (
+	"encoding/json"
 	"fmt"
 	"io"
 	"net"
@@ -46,6 +47,48 @@ func (c *awsCollector) isCollector() bool {
 	return c.cfg.StoreInterval != 0
 }
 
+func (c *awsCollector) hasReportsFromLive(ctx context.Context, userid string) (bool, error) {
+	span, ctx := opentracing.StartSpanFromContext(ctx, "hasReportsFromLive")
+	defer span.Finish()
+	if c.isCollector() {
+		e, found := c.pending.Load(userid)
+		if !found {
+			return false, nil
+		}
+		entry := e.(*pendingEntry)
+		entry.Lock()
+		defer entry.Unlock()
+		if entry.report != nil {
+			return true, nil
+		}
+		for _, v := range entry.older {
+			if v != nil {
+				return true, nil
+			}
+		}
+		return false, nil
+	}
+	// We are a querier: ask each collector if it has any
+	// (serially, since we will bail out on the first one that has reports)
+	addrs := resolve(c.cfg.CollectorAddr)
+	for _, addr := range addrs {
+		body, err := oneCall(ctx, addr, "/api/probes?sparse=true", userid)
+		if err != nil {
+			return false, err
+		}
+		var hasReports bool
+		decoder := json.NewDecoder(body)
+		if err := decoder.Decode(&hasReports); err != nil {
+			log.Errorf("Error encoding response: %v", err)
+		}
+		body.Close()
+		if hasReports {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
 func (c *awsCollector) reportsFromLive(ctx context.Context, userid string) ([]report.Report, error) {
 	span, ctx := opentracing.StartSpanFromContext(ctx, "reportsFromLive")
 	defer span.Finish()