BackendSrv: Support streaming chunked responses (#98691)

pkg/tsdb/grafana-testdata-datasource/resource_handler.go

@@ -4,13 +4,13 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math/rand"
 	"net/http"
 	"sort"
 	"strconv"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
-
 	"github.com/grafana/grafana/pkg/tsdb/grafana-testdata-datasource/kinds"
 )
 
@@ -78,34 +78,89 @@ func (s *Service) testStreamHandler(rw http.ResponseWriter, req *http.Request) {
 	ctxLogger := s.logger.FromContext(req.Context())
 	ctxLogger.Debug("Received resource call", "url", req.URL.String(), "method", req.Method)
 
+	header := rw.Header()
+	header.Set("Cache-Control", "no-store")
+	header.Set("X-Content-Type-Options", "nosniff")
+	header.Set("Content-Type", "text/plain")
+
+	writeError := func(code int, message string) {
+		rw.WriteHeader(code)
+		_, _ = rw.Write([]byte(message))
+	}
+
 	if req.Method != http.MethodGet {
+		writeError(http.StatusMethodNotAllowed, "only supports get")
 		return
 	}
 
+	var err error
+	query := req.URL.Query()
 	count := 10
-	countstr := req.URL.Query().Get("count")
-	if countstr != "" {
-		if i, err := strconv.Atoi(countstr); err == nil {
-			count = i
+	if query.Has("count") {
+		count, err = strconv.Atoi(query.Get("count"))
+		if err != nil {
+			writeError(http.StatusBadRequest, "invalid count value")
+			return
 		}
 	}
 
-	sleep := req.URL.Query().Get("sleep")
-	sleepDuration, err := time.ParseDuration(sleep)
-	if err != nil {
-		sleepDuration = time.Millisecond
+	start := 1
+	if query.Has("start") {
+		start, err = strconv.Atoi(query.Get("start"))
+		if err != nil {
+			writeError(http.StatusBadRequest, "invalid start value")
+			return
+		}
 	}
 
-	rw.Header().Set("Content-Type", "text/plain")
+	flush := 100 // flush 100% of the time
+	if query.Has("flush") {
+		flush, err = strconv.Atoi(query.Get("flush"))
+		if err != nil {
+			writeError(http.StatusBadRequest, "invalid flush value")
+			return
+		}
+		if flush > 100 || flush < 0 {
+			writeError(http.StatusBadRequest, "expecting flush between 0-100")
+			return
+		}
+	}
+
+	speed := time.Millisecond * 10
+	if query.Has("speed") {
+		speed, err = time.ParseDuration(query.Get("speed"))
+		if err != nil {
+			writeError(http.StatusBadRequest, "invalid speed")
+			return
+		}
+	}
+
+	line := func(i int) string {
+		return fmt.Sprintf("Message #%d", i)
+	}
+	switch query.Get("format") {
+	case "json":
+		line = func(i int) string {
+			return fmt.Sprintf(`{"message": %d, "value": %.3f, "time": %d}`, i, rand.Float64(), time.Now().UnixMilli())
+		}
+	case "influx":
+		line = func(i int) string {
+			val := rand.Float64()
+			return fmt.Sprintf("measurement,tag1=value1,tag2=value2 message=%d,value=%.3f %d", i, val, time.Now().UnixMilli())
+		}
+	}
 	rw.WriteHeader(http.StatusOK)
 
-	for i := 1; i <= count; i++ {
-		if _, err := io.WriteString(rw, fmt.Sprintf("Message #%d", i)); err != nil {
+	for i := start; i <= count; i++ {
+		if _, err := io.WriteString(rw, line(i)+"\n"); err != nil {
 			ctxLogger.Error("Failed to write response", "error", err)
 			return
 		}
-		rw.(http.Flusher).Flush()
-		time.Sleep(sleepDuration)
+		// This may send multiple lines in one chunk
+		if flush > rand.Intn(100) {
+			rw.(http.Flusher).Flush()
+		}
+		time.Sleep(speed)
 	}
 }