fix(loki-stream): reconnect when websocket is unexpectedly closed

go.mod

@@ -16,6 +16,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2
 	github.com/bep/debounce v1.2.1
 	github.com/cert-manager/cert-manager v1.20.2
+	github.com/coder/websocket v1.8.14
 	github.com/docker/cli v29.5.2+incompatible
 	github.com/friendsofgo/errors v0.9.2
 	github.com/gliderlabs/ssh v0.3.8
@@ -37,7 +38,6 @@ require (
 	github.com/prometheus/common v0.68.0
 	github.com/regclient/regclient v0.11.5
 	github.com/samber/lo v1.53.0
-	github.com/shiguredo/websocket v1.6.1
 	github.com/sourcegraph/conc v0.3.1-0.20240121214520-5f936abd7ae8
 	github.com/spf13/cobra v1.10.2
 	github.com/spf13/pflag v1.0.10

go.sum

@@ -97,6 +97,8 @@ github.com/cloudflare/circl v1.6.3 h1:9GPOhQGF9MCYUeXyMYlqTR6a5gTrgR/fBLXvUgtVcg
 github.com/cloudflare/circl v1.6.3/go.mod h1:2eXP6Qfat4O/Yhh8BznvKnJ+uzEoTQ6jVKJRn81BiS4=
 github.com/codahale/rfc6979 v0.0.0-20141003034818-6a90f24967eb h1:EDmT6Q9Zs+SbUoc7Ik9EfrFqcylYqgPZ9ANSbTAntnE=
 github.com/codahale/rfc6979 v0.0.0-20141003034818-6a90f24967eb/go.mod h1:ZjrT6AXHbDs86ZSdt/osfBi5qfexBrKUdONk989Wnk4=
+github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g=
+github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg=
 github.com/containerd/cgroups/v3 v3.1.3 h1:eUNflyMddm18+yrDmZPn3jI7C5hJ9ahABE5q6dyLYXQ=
 github.com/containerd/cgroups/v3 v3.1.3/go.mod h1:PKZ2AcWmSBsY/tJUVhtS/rluX0b1uq1GmPO1ElCmbOw=
 github.com/containerd/console v1.0.5 h1:R0ymNeydRqH2DmakFNdmjR2k0t7UPuiOV/N/27/qqsc=
@@ -443,8 +445,6 @@ github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3 h1:n661drycOFuPLCN
 github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=
 github.com/shibumi/go-pathspec v1.3.0 h1:QUyMZhFo0Md5B8zV8x2tesohbb5kfbpTi9rBnKh5dkI=
 github.com/shibumi/go-pathspec v1.3.0/go.mod h1:Xutfslp817l2I1cZvgcfeMQJG5QnU2lh5tVaaMCl3jE=
-github.com/shiguredo/websocket v1.6.1 h1:xGZ5LmGjQLfGaCcxZrI2/z0en24eJ3VunnNK1RmJdzg=
-github.com/shiguredo/websocket v1.6.1/go.mod h1:qUnxxJOWcK8Q7Q+o31UucO9HJQuhcygUYMuLRa88XTw=
 github.com/shirou/gopsutil/v4 v4.26.3 h1:2ESdQt90yU3oXF/CdOlRCJxrP+Am1aBYubTMTfxJ1qc=
 github.com/shirou/gopsutil/v4 v4.26.3/go.mod h1:LZ6ewCSkBqUpvSOf+LsTGnRinC6iaNUNMGBtDkJBaLQ=
 github.com/sigstore/sigstore v1.10.5 h1:KqrOjDhNOVY+uOzQFat2FrGLClPPCb3uz8pK3wuI+ow=

pkg/infrastructure/log/loki/streamer.go

@@ -5,14 +5,15 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"iter"
 	"log/slog"
 	"net/http"
 	"net/url"
 	"text/template"
 	"time"
 
+	"github.com/coder/websocket"
 	"github.com/friendsofgo/errors"
-	"github.com/shiguredo/websocket"
 
 	"github.com/traPtitech/neoshowcase/pkg/domain"
 )
@@ -122,46 +123,73 @@ func (l *lokiStreamer) Stream(ctx context.Context, app *domain.Application, begi
 	if err != nil {
 		return nil, errors.Wrap(err, "templating logQL")
 	}
-	q := make(url.Values)
-	q.Set("query", logQL)
-	q.Set("limit", "100")
-	// ensure start time is not in the future to prevent Loki API errors
-	start := min(time.Now().UnixNano(), begin.UnixNano())
-	q.Set("start", fmt.Sprintf("%d", start))
-
-	conn, _, err := websocket.DefaultDialer.DialContext(ctx, l.streamEndpoint()+"?"+q.Encode(), nil)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to dial to stream ws endpoint")
-	}
 
 	ch := make(chan *domain.ContainerLog, 100)
 
 	ctx, cancel := context.WithCancel(ctx)
 	go func() {
-		<-ctx.Done()
-		_ = conn.Close()
 		defer close(ch)
-	}()
-	go func() {
 		defer cancel()
-		defer slog.InfoContext(ctx, "closing loki websocket stream")
-		slog.InfoContext(ctx, "new loki websocket stream")
+
+		lastSeenTime := begin
 
 		for {
-			typ, b, err := conn.ReadMessage()
-			select { // check if context was cancelled
-			case <-ctx.Done():
+			logSeq, err := l.readStream(ctx, logQL, lastSeenTime)
+			if err != nil {
+				slog.ErrorContext(ctx, "failed to read log stream", "error", err)
 				return
-			default:
 			}
+			for l, err := range logSeq {
+				if errors.Is(err, context.Canceled) {
+					return
+				} else if err != nil {
+					slog.WarnContext(ctx, "unexpected error occurred while reading log stream", "error", err)
+					break // retry
+				}
+
+				switch {
+				case l.Time.After(lastSeenTime):
+					lastSeenTime = l.Time
+				case l.Time.Before(lastSeenTime):
+					continue
+				}
+				select {
+				case ch <- l:
+				case <-ctx.Done():
+					return
+				}
+			}
+		}
+	}()
+
+	return ch, nil
+}
+
+func (l *lokiStreamer) readStream(ctx context.Context, query string, start time.Time) (iter.Seq2[*domain.ContainerLog, error], error) {
+	q := make(url.Values)
+	q.Set("query", query)
+	q.Set("limit", "100")
+	q.Set("start", fmt.Sprintf("%d", start.UnixNano()))
+
+	conn, _, err := websocket.Dial(ctx, l.streamEndpoint()+"?"+q.Encode(), nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return func(yield func(*domain.ContainerLog, error) bool) {
+		defer conn.CloseNow()
+
+		for {
+			typ, b, err := conn.Read(ctx)
 			if err != nil {
-				slog.ErrorContext(ctx, "failed to read ws message", "error", err)
+				yield(nil, err)
 				return
 			}
+
 			switch typ {
-			case websocket.TextMessage:
+			case websocket.MessageText:
 				var res streamResponse
-				err = json.NewDecoder(bytes.NewReader(b)).Decode(&res)
+				err := json.Unmarshal(b, &res)
 				if err != nil {
 					slog.ErrorContext(ctx, "failed to decode ws message", "error", err)
 					continue // fail-safe
@@ -172,18 +200,13 @@ func (l *lokiStreamer) Stream(ctx context.Context, app *domain.Application, begi
 					continue // fail-safe
 				}
 				for _, l := range logs {
-					select {
-					case ch <- l:
-					default:
+					if !yield(l, nil) {
+						return
 					}
 				}
-			case websocket.BinaryMessage:
+			case websocket.MessageBinary:
 				// ignore
-			case websocket.CloseMessage:
-				return
 			}
 		}
-	}()
-
-	return ch, nil
+	}, nil
 }