Added support for local files to --log-output

cmd/root.go

@@ -246,7 +246,7 @@ func (c *rootCommand) rootCmdPersistentFlagSet() *pflag.FlagSet {
 	flags.BoolVarP(&c.commandFlags.quiet, "quiet", "q", false, "disable progress updates")
 	flags.BoolVar(&c.commandFlags.noColor, "no-color", false, "disable colored output")
 	flags.StringVar(&c.logOutput, "log-output", "stderr",
-		"change the output for k6 logs, possible values are stderr,stdout,none,loki[=host:port]")
+		"change the output for k6 logs, possible values are stderr,stdout,none,loki[=host:port],file[=./path.fileformat]")
 	flags.StringVar(&c.logFmt, "logformat", "", "log output format") // TODO rename to log-format and warn on old usage
 	flags.StringVarP(&c.commandFlags.address, "address", "a", "localhost:6565", "address for the api server")
 
@@ -279,27 +279,37 @@ func (c *rootCommand) setupLoggers() (<-chan struct{}, error) {
 	}
 
 	loggerForceColors := false // disable color by default
-	switch c.logOutput {
-	case "stderr":
+	switch line := c.logOutput; {
+	case line == "stderr":
 		loggerForceColors = !c.commandFlags.noColor && c.commandFlags.stderrTTY
 		c.logger.SetOutput(c.commandFlags.stderr)
-	case "stdout":
+	case line == "stdout":
 		loggerForceColors = !c.commandFlags.noColor && c.commandFlags.stdoutTTY
 		c.logger.SetOutput(c.commandFlags.stdout)
-	case "none":
+	case line == "none":
 		c.logger.SetOutput(ioutil.Discard)
-	default:
-		if !strings.HasPrefix(c.logOutput, "loki") {
-			return nil, fmt.Errorf("unsupported log output `%s`", c.logOutput)
-		}
+
+	case strings.HasPrefix(line, "loki"):
 		ch = make(chan struct{})
-		hook, err := log.LokiFromConfigLine(c.ctx, c.fallbackLogger, c.logOutput, ch)
+		hook, err := log.LokiFromConfigLine(c.ctx, c.fallbackLogger, line, ch)
 		if err != nil {
 			return nil, err
 		}
 		c.logger.AddHook(hook)
 		c.logger.SetOutput(ioutil.Discard) // don't output to anywhere else
 		c.logFmt = "raw"
+
+	case strings.HasPrefix(line, "file"):
+		hook, err := log.FileHookFromConfigLine(c.ctx, c.fallbackLogger, line)
+		if err != nil {
+			return nil, err
+		}
+
+		c.logger.AddHook(hook)
+		c.logger.SetOutput(ioutil.Discard)
+
+	default:
+		return nil, fmt.Errorf("unsupported log output `%s`", line)
 	}
 
 	switch c.logFmt {

log/file.go

@@ -0,0 +1,160 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2020 Load Impact
+ *
+ * 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 implements various logrus hooks.
+package log
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/sirupsen/logrus"
+)
+
+// fileHookBufferSize is a default size for the fileHook's loglines channel.
+const fileHookBufferSize = 100
+
+// fileHook is a hook to handle writing to local files.
+type fileHook struct {
+	fallbackLogger logrus.FieldLogger
+	loglines       chan []byte
+	path           string
+	w              io.WriteCloser
+	bw             *bufio.Writer
+	levels         []logrus.Level
+}
+
+// FileHookFromConfigLine returns new fileHook hook.
+func FileHookFromConfigLine(
+	ctx context.Context, fallbackLogger logrus.FieldLogger, line string,
+) (logrus.Hook, error) {
+	hook := &fileHook{
+		fallbackLogger: fallbackLogger,
+		levels:         logrus.AllLevels,
+	}
+
+	parts := strings.SplitN(line, "=", 2)
+	if parts[0] != "file" {
+		return nil, fmt.Errorf("logfile configuration should be in the form `file=path-to-local-file` but is `%s`", line)
+	}
+
+	if err := hook.parseArgs(line); err != nil {
+		return nil, err
+	}
+
+	if err := hook.openFile(); err != nil {
+		return nil, err
+	}
+
+	hook.loglines = hook.loop(ctx)
+
+	return hook, nil
+}
+
+func (h *fileHook) parseArgs(line string) error {
+	tokens, err := tokenize(line)
+	if err != nil {
+		return fmt.Errorf("error while parsing logfile configuration %w", err)
+	}
+
+	for _, token := range tokens {
+		switch token.key {
+		case "file":
+			if token.value == "" {
+				return fmt.Errorf("filepath must not be empty")
+			}
+			h.path = token.value
+		case "level":
+			h.levels, err = parseLevels(token.value)
+			if err != nil {
+				return err
+			}
+		default:
+			return fmt.Errorf("unknown logfile config key %s", token.key)
+		}
+	}
+
+	return nil
+}
+
+// openFile opens logfile and initializes writers.
+func (h *fileHook) openFile() error {
+	if _, err := os.Stat(filepath.Dir(h.path)); os.IsNotExist(err) {
+		return fmt.Errorf("provided directory '%s' does not exist", filepath.Dir(h.path))
+	}
+
+	file, err := os.OpenFile(h.path, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0o600)
+	if err != nil {
+		return fmt.Errorf("failed to open logfile %s: %w", h.path, err)
+	}
+
+	h.w = file
+	h.bw = bufio.NewWriter(file)
+
+	return nil
+}
+
+func (h *fileHook) loop(ctx context.Context) chan []byte {
+	loglines := make(chan []byte, fileHookBufferSize)
+
+	go func() {
+		for {
+			select {
+			case entry := <-loglines:
+				if _, err := h.bw.Write(entry); err != nil {
+					h.fallbackLogger.Errorf("failed to write a log message to a logfile: %w", err)
+				}
+			case <-ctx.Done():
+				if err := h.bw.Flush(); err != nil {
+					h.fallbackLogger.Errorf("failed to flush buffer: %w", err)
+				}
+
+				if err := h.w.Close(); err != nil {
+					h.fallbackLogger.Errorf("failed to close logfile: %w", err)
+				}
+
+				return
+			}
+		}
+	}()
+
+	return loglines
+}
+
+// Fire writes the log file to defined path.
+func (h *fileHook) Fire(entry *logrus.Entry) error {
+	message, err := entry.Bytes()
+	if err != nil {
+		return fmt.Errorf("failed to get a log entry bytes: %w", err)
+	}
+
+	h.loglines <- message
+	return nil
+}
+
+// Levels returns configured log levels.
+func (h *fileHook) Levels() []logrus.Level {
+	return h.levels
+}

log/file_test.go

@@ -0,0 +1,168 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2020 Load Impact
+ *
+ * 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
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type nopCloser struct {
+	io.Writer
+	closed chan struct{}
+}
+
+func (nc *nopCloser) Close() error {
+	nc.closed <- struct{}{}
+	return nil
+}
+
+func TestFileHookFromConfigLine(t *testing.T) {
+	t.Parallel()
+
+	tests := [...]struct {
+		line       string
+		err        bool
+		errMessage string
+		res        fileHook
+	}{
+		{
+			line: "file",
+			err:  true,
+			res: fileHook{
+				levels: logrus.AllLevels,
+			},
+		},
+		{
+			line: fmt.Sprintf("file=%s/k6.log,level=info", os.TempDir()),
+			err:  false,
+			res: fileHook{
+				path:   fmt.Sprintf("%s/k6.log", os.TempDir()),
+				levels: logrus.AllLevels[:5],
+			},
+		},
+		{
+			line: "file=./",
+			err:  true,
+		},
+		{
+			line: "file=/a/c/",
+			err:  true,
+		},
+		{
+			line:       "file=,level=info",
+			err:        true,
+			errMessage: "filepath must not be empty",
+		},
+		{
+			line: "file=/tmp/k6.log,level=tea",
+			err:  true,
+		},
+		{
+			line: "file=/tmp/k6.log,unknown",
+			err:  true,
+		},
+		{
+			line: "file=/tmp/k6.log,level=",
+			err:  true,
+		},
+		{
+			line: "file=/tmp/k6.log,level=,",
+			err:  true,
+		},
+		{
+			line:       "file=/tmp/k6.log,unknown=something",
+			err:        true,
+			errMessage: "unknown logfile config key unknown",
+		},
+		{
+			line:       "unknown=something",
+			err:        true,
+			errMessage: "logfile configuration should be in the form `file=path-to-local-file` but is `unknown=something`",
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.line, func(t *testing.T) {
+			t.Parallel()
+
+			res, err := FileHookFromConfigLine(context.Background(), logrus.New(), test.line)
+
+			if test.err {
+				require.Error(t, err)
+
+				if test.errMessage != "" {
+					require.Equal(t, test.errMessage, err.Error())
+				}
+
+				return
+			}
+
+			require.NoError(t, err)
+			assert.NotNil(t, res.(*fileHook).w)
+		})
+	}
+}
+
+func TestFileHookFire(t *testing.T) {
+	t.Parallel()
+
+	var buffer bytes.Buffer
+	nc := &nopCloser{
+		Writer: &buffer,
+		closed: make(chan struct{}),
+	}
+
+	hook := &fileHook{
+		loglines: make(chan []byte),
+		w:        nc,
+		bw:       bufio.NewWriter(nc),
+		levels:   logrus.AllLevels,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	hook.loglines = hook.loop(ctx)
+
+	logger := logrus.New()
+	logger.AddHook(hook)
+	logger.SetOutput(io.Discard)
+
+	logger.Info("example log line")
+
+	time.Sleep(10 * time.Millisecond)
+
+	cancel()
+	<-nc.closed
+
+	assert.Contains(t, buffer.String(), "example log line")
+}