fmsg: implement suspend in writer

This removes the requirement to call fmsg.Exit on every exit path, and enables direct use of the "log" package. However, fmsg.BeforeExit is still encouraged when possible to catch exit on suspended output.

Signed-off-by: Ophestra <cat@gensokyo.uk>

internal/fmsg/fmsg.go

@@ -2,39 +2,85 @@
 package fmsg
 
 import (
+	"bytes"
+	"io"
 	"log"
 	"os"
+	"sync"
+	"sync/atomic"
+	"syscall"
 )
 
-var std = log.New(os.Stderr, "fortify: ", 0)
+const (
+	bufSize    = 4 * 1024
+	bufSizeMax = 16 * 1024 * 1024
+)
 
-func SetPrefix(prefix string) {
-	prefix += ": "
-	std.SetPrefix(prefix)
-	std.SetPrefix(prefix)
+var o = &suspendable{w: os.Stderr}
+
+// Prepare configures the system logger for [Suspend] and [Resume] to take effect.
+func Prepare(prefix string) { log.SetPrefix(prefix + ": "); log.SetFlags(0); log.SetOutput(o) }
+
+type suspendable struct {
+	w io.Writer
+	s atomic.Bool
+
+	buf     bytes.Buffer
+	bufOnce sync.Once
+	bufMu   sync.Mutex
+	dropped int
 }
 
-func Print(v ...any) {
-	dequeueOnce.Do(dequeue)
-	queue(dPrint(v))
+func (s *suspendable) Write(p []byte) (n int, err error) {
+	if !s.s.Load() {
+		return s.w.Write(p)
+	}
+	s.bufOnce.Do(func() { s.prepareBuf() })
+
+	s.bufMu.Lock()
+	defer s.bufMu.Unlock()
+
+	if l := len(p); s.buf.Len()+l > bufSizeMax {
+		s.dropped += l
+		return 0, syscall.ENOMEM
+	}
+	return s.buf.Write(p)
 }
 
-func Printf(format string, v ...any) {
-	dequeueOnce.Do(dequeue)
-	queue(&dPrintf{format, v})
+func (s *suspendable) prepareBuf()   { s.buf.Grow(bufSize) }
+func (s *suspendable) Suspend() bool { return o.s.CompareAndSwap(false, true) }
+func (s *suspendable) Resume() (resumed bool, dropped uintptr, n int64, err error) {
+	if o.s.CompareAndSwap(true, false) {
+		o.bufMu.Lock()
+		defer o.bufMu.Unlock()
+
+		resumed = true
+		dropped = uintptr(o.dropped)
+
+		o.dropped = 0
+		n, err = io.Copy(s.w, &s.buf)
+		s.buf = bytes.Buffer{}
+		s.prepareBuf()
+	}
+	return
 }
 
-func Println(v ...any) {
-	dequeueOnce.Do(dequeue)
-	queue(dPrintln(v))
+func Suspend() bool { return o.Suspend() }
+func Resume() bool {
+	resumed, dropped, _, err := o.Resume()
+	if err != nil {
+		// probably going to result in an error as well,
+		// so this call is as good as unreachable
+		log.Printf("cannot dump buffer on resume: %v", err)
+	}
+	if resumed && dropped > 0 {
+		log.Fatalf("dropped %d bytes while output is suspended", dropped)
+	}
+	return resumed
 }
 
-func Fatal(v ...any) {
-	Print(v...)
-	Exit(1)
-}
-
-func Fatalf(format string, v ...any) {
-	Printf(format, v...)
-	Exit(1)
+func BeforeExit() {
+	if Resume() {
+		log.Printf("beforeExit reached on suspended output")
+	}
 }