From e599b5583dd821ed26bdfc01f7f67d4a317b9843 Mon Sep 17 00:00:00 2001
From: Ophestra <cat@gensokyo.uk>
Date: Sun, 16 Feb 2025 17:26:09 +0900
Subject: 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 | 88 +++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 67 insertions(+), 21 deletions(-)

(limited to 'internal/fmsg/fmsg.go')

diff --git a/internal/fmsg/fmsg.go b/internal/fmsg/fmsg.go
index d4027204..48538380 100644
--- a/internal/fmsg/fmsg.go
+++ b/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
 
-func Print(v ...any) {
-	dequeueOnce.Do(dequeue)
-	queue(dPrint(v))
+	buf     bytes.Buffer
+	bufOnce sync.Once
+	bufMu   sync.Mutex
+	dropped int
 }
 
-func Printf(format string, v ...any) {
-	dequeueOnce.Do(dequeue)
-	queue(&dPrintf{format, 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 Println(v ...any) {
-	dequeueOnce.Do(dequeue)
-	queue(dPrintln(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 Fatal(v ...any) {
-	Print(v...)
-	Exit(1)
+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 Fatalf(format string, v ...any) {
-	Printf(format, v...)
-	Exit(1)
+func BeforeExit() {
+	if Resume() {
+		log.Printf("beforeExit reached on suspended output")
+	}
 }
-- 
cgit v1.3.1