Implemented a wait function to get all queued messages sent before exiting

lib.go

@@ -25,13 +25,15 @@ type VictoriaHandler struct {
 	Instance    string
 	URL         string
 
-	handler     slog.Handler
+	waitCalled      chan bool
-	filecounter int
+	queueLoopExited chan bool
-	logdir      string
+	handler         slog.Handler
-	logQueue    chan []byte
+	filecounter     int
-	run         string
+	logdir          string
-	log         *slog.Logger
+	logQueue        chan []byte
-	httpClient  http.Client
+	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)
+	msgLoop:
-		if err != nil {
+		for {
-			vl.logToFile(msg)
+			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) { // {{{