From 08eb271a4ca257aaff5024e6bbd844a0182f0f78 Mon Sep 17 00:00:00 2001
From: kim <89579420+NyaaaWhatsUpDoc@users.noreply.github.com>
Date: Sat, 7 May 2022 16:46:59 +0100
Subject: [PATCH] [performance] improved logrus output switching performance
 (#544)

* improved logrus output switching performance

Signed-off-by: kim <grufwub@gmail.com>

* add license to test file

Signed-off-by: kim <grufwub@gmail.com>
---
 internal/log/log.go      | 45 ++++++++++++++++++--------
 internal/log/log_test.go | 69 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 13 deletions(-)
 create mode 100644 internal/log/log_test.go

diff --git a/internal/log/log.go b/internal/log/log.go
index d2114692c..f5ce7ec24 100644
--- a/internal/log/log.go
+++ b/internal/log/log.go
@@ -20,9 +20,9 @@
 
 import (
 	"bytes"
-	"os"
-
+	"io"
 	"log/syslog"
+	"os"
 
 	"github.com/sirupsen/logrus"
 	lSyslog "github.com/sirupsen/logrus/hooks/syslog"
@@ -34,12 +34,13 @@
 // log level from the viper store, or using a default if the level
 // has not been set in viper.
 //
-// It also sets the output to log.outputSplitter,
+// It also sets the output to log.SplitErrOutputs(...)
 // so you get error logs on stderr and normal logs on stdout.
 //
 // If syslog settings are also in viper, then Syslog will be initialized as well.
 func Initialize() error {
-	logrus.SetOutput(&outputSplitter{})
+	out := SplitErrOutputs(os.Stdout, os.Stderr)
+	logrus.SetOutput(out)
 
 	logrus.SetFormatter(&logrus.TextFormatter{
 		DisableColors: true,
@@ -79,14 +80,32 @@ func Initialize() error {
 	return nil
 }
 
-// outputSplitter implements the io.Writer interface for use with Logrus, and simply
-// splits logs between stdout and stderr depending on their severity.
-// See: https://github.com/sirupsen/logrus/issues/403#issuecomment-346437512
-type outputSplitter struct{}
-
-func (splitter *outputSplitter) Write(p []byte) (n int, err error) {
-	if bytes.Contains(p, []byte("level=error")) {
-		return os.Stderr.Write(p)
+// SplitErrOutputs returns an OutputSplitFunc that splits output to either one of
+// two given outputs depending on whether the level is "error","fatal","panic".
+func SplitErrOutputs(out, err io.Writer) OutputSplitFunc {
+	return func(lvl []byte) io.Writer {
+		switch string(lvl) /* convert to str for compare is no-alloc */ {
+		case "error", "fatal", "panic":
+			return err
+		default:
+			return out
+		}
 	}
-	return os.Stdout.Write(p)
+}
+
+// OutputSplitFunc implements the io.Writer interface for use with Logrus, and simply
+// splits logs between stdout and stderr depending on their severity.
+type OutputSplitFunc func(lvl []byte) io.Writer
+
+var levelBytes = []byte("level=")
+
+func (fn OutputSplitFunc) Write(b []byte) (int, error) {
+	var lvl []byte
+	if i := bytes.Index(b, levelBytes); i >= 0 {
+		blvl := b[i+len(levelBytes):]
+		if i := bytes.IndexByte(blvl, ' '); i >= 0 {
+			lvl = blvl[:i]
+		}
+	}
+	return fn(lvl).Write(b)
 }
diff --git a/internal/log/log_test.go b/internal/log/log_test.go
new file mode 100644
index 000000000..01aae73ae
--- /dev/null
+++ b/internal/log/log_test.go
@@ -0,0 +1,69 @@
+/*
+   GoToSocial
+   Copyright (C) 2021-2022 GoToSocial Authors admin@gotosocial.org
+
+   This program is free software: you can redistribute it and/or modify
+   it under the terms of the GNU Affero General Public License as published by
+   the Free Software Foundation, either version 3 of the License, or
+   (at your option) any later version.
+
+   This program is distributed in the hope that it will be useful,
+   but WITHOUT ANY WARRANTY; without even the implied warranty of
+   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+   GNU Affero General Public License for more details.
+
+   You should have received a copy of the GNU Affero General Public License
+   along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package log_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/sirupsen/logrus"
+	"github.com/superseriousbusiness/gotosocial/internal/log"
+)
+
+func TestOutputSplitFunc(t *testing.T) {
+	var outbuf, errbuf bytes.Buffer
+
+	out := log.SplitErrOutputs(&outbuf, &errbuf)
+
+	log := logrus.New()
+	log.SetOutput(out)
+	log.SetLevel(logrus.TraceLevel)
+
+	for _, lvl := range logrus.AllLevels {
+		func() {
+			defer func() { recover() }()
+			log.Log(lvl, "hello world")
+		}()
+
+		t.Logf("outbuf=%q errbuf=%q", outbuf.String(), errbuf.String())
+
+		switch lvl {
+		case logrus.PanicLevel:
+			if outbuf.Len() > 0 || errbuf.Len() == 0 {
+				t.Error("expected panic to log to OutputSplitter.Err")
+			}
+		case logrus.FatalLevel:
+			if outbuf.Len() > 0 || errbuf.Len() == 0 {
+				t.Error("expected fatal to log to OutputSplitter.Err")
+			}
+		case logrus.ErrorLevel:
+			if outbuf.Len() > 0 || errbuf.Len() == 0 {
+				t.Error("expected error to log to OutputSplitter.Err")
+			}
+		default:
+			if outbuf.Len() == 0 || errbuf.Len() > 0 {
+				t.Errorf("expected %s to log to OutputSplitter.Out", lvl)
+			}
+		}
+
+		// Reset buffers
+		outbuf.Reset()
+		errbuf.Reset()
+	}
+}