fix(observability-proxy): tell a slow store apart and bound tail memory

A request that hits the proxy's own deadline was answered like an
unreachable store, so a heavy query looked like an outage; it now
answers 504, and a client that went away is no longer logged as a
store failure. A tail line is read through a fixed buffer, so a single
huge log line cannot make one stream hold unbounded memory.

Assisted-by: LLM
Signed-off-by: ZverGuy <maximbel2003@gmail.com>

internal/observabilityproxy/server.go

@@ -31,6 +31,10 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 )
 
+// tailChunkBytes bounds the memory one tail holds: a longer line is passed
+// on in pieces and flushed once its end arrives.
+const tailChunkBytes = 64 << 10
+
 // maxFormBytes bounds a POST body; LogsQL queries are a few KiB at most.
 const maxFormBytes = 1 << 20
 
@@ -194,7 +198,13 @@ func (s *Server) serve(w http.ResponseWriter, r *http.Request, route Route) *HTT
 	req.Header.Set("AccountID", "0")
 	req.Header.Set("ProjectID", "0")
 	resp, err := s.client.Do(req)
-	if err != nil {
+	switch {
+	case err == nil:
+	case r.Context().Err() != nil:
+		return errorf(http.StatusBadGateway, "client went away")
+	case errors.Is(err, context.DeadlineExceeded):
+		return errorf(http.StatusGatewayTimeout, "log store did not answer within %s", s.cfg.RequestTimeout)
+	default:
 		slog.Warn("log store unreachable", "store", store, "error", err)
 		return errorf(http.StatusBadGateway, "log store unreachable")
 	}
@@ -249,21 +259,24 @@ func (s *Server) streamLines(w http.ResponseWriter, rc *http.ResponseController,
 		idle = time.AfterFunc(s.cfg.TailIdleTimeout, cancel)
 		defer idle.Stop()
 	}
-	rd := bufio.NewReader(body)
+	rd := bufio.NewReaderSize(body, tailChunkBytes)
 	for {
-		line, err := rd.ReadBytes('\n')
-		if len(line) > 0 {
+		chunk, err := rd.ReadSlice('\n')
+		if len(chunk) > 0 {
 			if idle != nil {
 				idle.Reset(s.cfg.TailIdleTimeout)
 			}
-			if _, werr := w.Write(line); werr != nil {
+			if _, werr := w.Write(chunk); werr != nil {
 				return
 			}
+			if errors.Is(err, bufio.ErrBufferFull) {
+				continue
+			}
 			if ferr := rc.Flush(); ferr != nil {
 				return
 			}
 		}
-		if err != nil {
+		if err != nil && !errors.Is(err, bufio.ErrBufferFull) {
 			return
 		}
 	}

internal/observabilityproxy/server_test.go

@@ -366,3 +366,30 @@ func TestServerTailCapsAndIdle(t *testing.T) {
 		time.Sleep(10 * time.Millisecond)
 	}
 }
+
+func TestServerUpstreamDeadline(t *testing.T) {
+	release := make(chan struct{})
+	t.Cleanup(func() { close(release) })
+	h := newHarness(t, Config{RequestTimeout: 50 * time.Millisecond}, func(_ http.ResponseWriter, r *http.Request) {
+		select {
+		case <-release:
+		case <-r.Context().Done():
+		}
+	})
+	if got := h.status(http.MethodGet, logsRoot+"query", "good"); got != http.StatusGatewayTimeout {
+		t.Fatalf("slow upstream answered %d, want 504", got)
+	}
+}
+
+func TestServerTailLongLine(t *testing.T) {
+	long := `{"_msg":"` + strings.Repeat("x", 3*tailChunkBytes) + `"}`
+	h := newHarness(t, Config{}, func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = io.WriteString(w, long+"\n"+`{"_msg":"short"}`+"\n")
+	})
+	resp := h.do(http.MethodGet, logsRoot+"tail", "good", nil, nil)
+	body, _ := io.ReadAll(resp.Body)
+	_ = resp.Body.Close()
+	if string(body) != long+"\n"+`{"_msg":"short"}`+"\n" {
+		t.Fatalf("long line not streamed intact: got %d bytes", len(body))
+	}
+}