From d7df24c9992bd29da1d775c45cb73eb7892e5f78 Mon Sep 17 00:00:00 2001
From: Ophestra Umiker <cat@ophivana.moe>
Date: Mon, 4 Nov 2024 12:56:19 +0900
Subject: [PATCH] fmsg: drop messages when msgbuf is full during withhold

Logging functions are not expected to block. This change fixes multiple hangs where more than 64 messages are produced during withhold.

Signed-off-by: Ophestra Umiker <cat@ophivana.moe>
---
 internal/fmsg/defer.go | 23 +++++++++++++++++++++++
 internal/fmsg/fmsg.go  |  9 +++------
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/internal/fmsg/defer.go b/internal/fmsg/defer.go
index deef58f..ed4a773 100644
--- a/internal/fmsg/defer.go
+++ b/internal/fmsg/defer.go
@@ -8,6 +8,7 @@ import (
 
 var (
 	wstate   atomic.Bool
+	dropped  atomic.Uint64
 	withhold = make(chan struct{}, 1)
 	msgbuf   = make(chan dOp, 64) // these ops are tiny so a large buffer is allocated for withholding output
 
@@ -29,6 +30,25 @@ func dequeue() {
 	}()
 }
 
+// queue submits ops to msgbuf but drops messages
+// when the buffer is full and dequeue is withholding
+func queue(op dOp) {
+	select {
+	case msgbuf <- op:
+		queueSync.Add(1)
+	default:
+		// send the op anyway if not withholding
+		// as dequeue will get to it eventually
+		if !wstate.Load() {
+			queueSync.Add(1)
+			msgbuf <- op
+		} else {
+			// increment dropped message count
+			dropped.Add(1)
+		}
+	}
+}
+
 type dOp interface{ Do() }
 
 func Exit(code int) {
@@ -47,6 +67,9 @@ func Resume() {
 	dequeueOnce.Do(dequeue)
 	if wstate.CompareAndSwap(true, false) {
 		withhold <- struct{}{}
+		if d := dropped.Swap(0); d != 0 {
+			Printf("dropped %d messages during withhold", d)
+		}
 	}
 }
 
diff --git a/internal/fmsg/fmsg.go b/internal/fmsg/fmsg.go
index dbe590b..d402720 100644
--- a/internal/fmsg/fmsg.go
+++ b/internal/fmsg/fmsg.go
@@ -16,20 +16,17 @@ func SetPrefix(prefix string) {
 
 func Print(v ...any) {
 	dequeueOnce.Do(dequeue)
-	queueSync.Add(1)
-	msgbuf <- dPrint(v)
+	queue(dPrint(v))
 }
 
 func Printf(format string, v ...any) {
 	dequeueOnce.Do(dequeue)
-	queueSync.Add(1)
-	msgbuf <- &dPrintf{format, v}
+	queue(&dPrintf{format, v})
 }
 
 func Println(v ...any) {
 	dequeueOnce.Do(dequeue)
-	queueSync.Add(1)
-	msgbuf <- dPrintln(v)
+	queue(dPrintln(v))
 }
 
 func Fatal(v ...any) {