From d45003d1a362eaf9582ccecc7f0292e7dcb404bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Magnus=20=C3=85hall?= <magnus.ahall@gibon.se>
Date: Thu, 19 Mar 2026 10:16:25 +0100
Subject: [PATCH] Implemented a wait function to get all queued messages sent
 before exiting

---
 lib.go | 61 +++++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 48 insertions(+), 13 deletions(-)

diff --git a/lib.go b/lib.go
index 5139676..19fb6d7 100644
--- a/lib.go
+++ b/lib.go
@@ -25,13 +25,15 @@ type VictoriaHandler struct {
 	Instance    string
 	URL         string
 
-	handler     slog.Handler
-	filecounter int
-	logdir      string
-	logQueue    chan []byte
-	run         string
-	log         *slog.Logger
-	httpClient  http.Client
+	waitCalled      chan bool
+	queueLoopExited chan bool
+	handler         slog.Handler
+	filecounter     int
+	logdir          string
+	logQueue        chan []byte
+	run             string
+	log             *slog.Logger
+	httpClient      http.Client
 }
 
 type LogData struct {
@@ -49,6 +51,8 @@ func New(w io.Writer, opts slog.HandlerOptions, logdir, url, application, system
 	vl.logdir = logdir
 
 	vl.handler = slog.NewJSONHandler(w, &opts)
+	vl.waitCalled = make(chan bool)
+	vl.queueLoopExited = make(chan bool)
 
 	// logdir is used to spool logfiles in case of
 	err = os.MkdirAll(logdir, 0700)
@@ -71,10 +75,14 @@ func New(w io.Writer, opts slog.HandlerOptions, logdir, url, application, system
 } // }}}
 
 // Handle is overridden to enable central logging.
-func (vl VictoriaHandler) Handle(ctx context.Context, record slog.Record) error {// {{{
+func (vl VictoriaHandler) Handle(ctx context.Context, record slog.Record) error { // {{{
 	vl.queue(record)
 	return vl.handler.Handle(ctx, record)
-}// }}}
+} // }}}
+func (vl VictoriaHandler) Wait() { // {{{
+	vl.waitCalled <- true
+	<- vl.queueLoopExited
+} // }}}
 
 func (vl *VictoriaHandler) diskHandler() { // {{{
 	for {
@@ -113,13 +121,40 @@ func (vl *VictoriaHandler) diskHandler() { // {{{
 } // }}}
 func (vl *VictoriaHandler) queueHandler() { // {{{
 	var err error
+
+queueLoop:
 	for {
-		msg := <-vl.logQueue
-		err = vl.send(msg)
-		if err != nil {
-			vl.logToFile(msg)
+
+	msgLoop:
+		for {
+			select {
+			case msg := <-vl.logQueue:
+				err = vl.send(msg)
+				if err != nil {
+					vl.logToFile(msg)
+				}
+
+			// This can be selected while there still are messages in logQueue.
+			case <-vl.waitCalled:
+				break msgLoop
+			}
+		}
+
+		// Another turn on log queue to gather up all remaining messages.
+		for {
+			select {
+			case msg := <-vl.logQueue:
+				err = vl.send(msg)
+				if err != nil {
+					vl.logToFile(msg)
+				}
+			default:
+				break queueLoop
+			}
 		}
 	}
+
+	vl.queueLoopExited <- true
 } // }}}
 
 func (vl *VictoriaHandler) queue(rec slog.Record) { // {{{