NOISSUE - Cache and retry message sending (#222)

* cache and retry message sending

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* cache and retry message sending

Signed-off-by: SammyOina <sammyoina@gmail.com>

* remove safeconn

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* simplify retry

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* debug disconnect

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* remove debug

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* simplify

Signed-off-by: SammyOina <sammyoina@gmail.com>

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>
Signed-off-by: SammyOina <sammyoina@gmail.com>

internal/logger/protohandler.go

@@ -6,32 +6,44 @@ import (
 	"context"
 	"io"
 	"log/slog"
+	"sync"
+	"time"
 
 	"github.com/ultravioletrs/cocos/pkg/manager"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
+const retryInterval = 5 * time.Second
+
 var _ slog.Handler = (*handler)(nil)
 
 type handler struct {
-	opts  slog.HandlerOptions
-	w     io.Writer
-	cmpID string
+	opts           slog.HandlerOptions
+	w              io.Writer
+	cmpID          string
+	cachedMessages [][]byte
+	mutex          sync.Mutex
+	stopRetry      chan struct{}
 }
 
-func NewProtoHandler(w io.Writer, opts *slog.HandlerOptions, cmpID string) slog.Handler {
+func NewProtoHandler(conn io.Writer, opts *slog.HandlerOptions, cmpID string) slog.Handler {
 	if opts == nil {
 		opts = &slog.HandlerOptions{}
 	}
-	return &handler{
-		opts:  *opts,
-		w:     w,
-		cmpID: cmpID,
+	h := &handler{
+		opts:           *opts,
+		w:              conn,
+		cmpID:          cmpID,
+		cachedMessages: make([][]byte, 0),
+		stopRetry:      make(chan struct{}),
 	}
+
+	go h.periodicRetry()
+
+	return h
 }
 
-// Enabled implements slog.Handler.
 func (h *handler) Enabled(_ context.Context, l slog.Level) bool {
 	minLevel := slog.LevelInfo
 	if h.opts.Level != nil {
@@ -40,13 +52,11 @@ func (h *handler) Enabled(_ context.Context, l slog.Level) bool {
 	return l >= minLevel
 }
 
-// Handle implements slog.Handler.
 func (h *handler) Handle(_ context.Context, r slog.Record) error {
 	message := r.Message
 	timestamp := timestamppb.New(r.Time)
 	level := r.Level.String()
 
-	// Calculate the number of chunks
 	chunkSize := 500
 	numChunks := (len(message) + chunkSize - 1) / chunkSize
 
@@ -57,10 +67,8 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
 			end = len(message)
 		}
 
-		// Create a chunk of the message
 		chunk := message[start:end]
 
-		// Create the agent log with the chunk
 		agentLog := manager.ClientStreamMessage{
 			Message: &manager.ClientStreamMessage_AgentLog{
 				AgentLog: &manager.AgentLog{
@@ -72,27 +80,57 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
 			},
 		}
 
-		// Marshal the chunk to protobuf
 		b, err := proto.Marshal(&agentLog)
 		if err != nil {
 			return err
 		}
 
-		// Write the chunk to the writer
-		if _, err := h.w.Write(b); err != nil {
-			return err
+		h.mutex.Lock()
+		_, err = h.w.Write(b)
+		if err != nil {
+			h.cachedMessages = append(h.cachedMessages, b)
 		}
+		h.mutex.Unlock()
 	}
 
 	return nil
 }
 
-// WithAttrs implements slog.Handler.
-func (*handler) WithAttrs(attrs []slog.Attr) slog.Handler {
+func (h *handler) periodicRetry() {
+	ticker := time.NewTicker(retryInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			h.retrySendCachedMessages()
+		case <-h.stopRetry:
+			return
+		}
+	}
+}
+
+func (h *handler) retrySendCachedMessages() {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+	tmp := [][]byte{}
+	for _, msg := range h.cachedMessages {
+		if _, err := h.w.Write(msg); err != nil {
+			tmp = append(tmp, msg)
+		}
+	}
+	h.cachedMessages = tmp
+}
+
+func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler {
 	panic("unimplemented")
 }
 
-// WithGroup implements slog.Handler.
-func (*handler) WithGroup(name string) slog.Handler {
+func (h *handler) WithGroup(name string) slog.Handler {
 	panic("unimplemented")
 }
+
+func (h *handler) Close() error {
+	close(h.stopRetry)
+	return nil
+}