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/defer.go

@@ -1,98 +0,0 @@
-package fmsg
-
-import (
-	"os"
-	"sync"
-	"sync/atomic"
-)
-
-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
-
-	dequeueOnce sync.Once
-	queueSync   sync.WaitGroup
-)
-
-func dequeue() {
-	go func() {
-		for {
-			select {
-			case op := <-msgbuf:
-				op.Do()
-				queueSync.Done()
-			case <-withhold:
-				<-withhold
-			}
-		}
-	}()
-}
-
-// queue submits ops to msgbuf but drops messages
-// when the buffer is full and dequeue is withholding
-func queue(op dOp) {
-	queueSync.Add(1)
-
-	select {
-	case msgbuf <- op:
-	default:
-		// send the op anyway if not withholding
-		// as dequeue will get to it eventually
-		if !wstate.Load() {
-			msgbuf <- op
-		} else {
-			queueSync.Done()
-			// increment dropped message count
-			dropped.Add(1)
-		}
-	}
-}
-
-type dOp interface{ Do() }
-
-func Exit(code int) {
-	Resume() // resume here to avoid deadlock
-	queueSync.Wait()
-	os.Exit(code)
-}
-
-func Suspend() {
-	dequeueOnce.Do(dequeue)
-	if wstate.CompareAndSwap(false, true) {
-		queueSync.Wait()
-		withhold <- struct{}{}
-	}
-}
-
-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)
-		}
-	}
-}
-
-type dPrint []any
-
-func (v dPrint) Do() {
-	std.Print(v...)
-}
-
-type dPrintf struct {
-	format string
-	v      []any
-}
-
-func (d *dPrintf) Do() {
-	std.Printf(d.format, d.v...)
-}
-
-type dPrintln []any
-
-func (v dPrintln) Do() {
-	std.Println(v...)
-}

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")
+	}
 }

internal/fmsg/verbose.go

@@ -1,25 +1,23 @@
 package fmsg
 
-import "sync/atomic"
+import (
+	"log"
+	"sync/atomic"
+)
 
 var verbose = new(atomic.Bool)
 
-func Verbose() bool {
-	return verbose.Load()
-}
+func Load() bool   { return verbose.Load() }
+func Store(v bool) { verbose.Store(v) }
 
-func SetVerbose(v bool) {
-	verbose.Store(v)
-}
-
-func VPrintf(format string, v ...any) {
+func Verbosef(format string, v ...any) {
 	if verbose.Load() {
-		Printf(format, v...)
+		log.Printf(format, v...)
 	}
 }
 
-func VPrintln(v ...any) {
+func Verbose(v ...any) {
 	if verbose.Load() {
-		Println(v...)
+		log.Println(v...)
 	}
 }