Wait for file log to drain before k6 exits

grafana · Mar 7, 2022 · fea17a1 · fea17a1
1 parent 5ab5c98
commit fea17a1
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -290,7 +290,7 @@ func (c *rootCommand) setupLoggers() (<-chan struct{}, error) {
 		c.logger.SetOutput(ioutil.Discard)
 
 	case strings.HasPrefix(line, "loki"):
-		ch = make(chan struct{})
+		ch = make(chan struct{}) // TODO: refactor, get it from the constructor
 		hook, err := log.LokiFromConfigLine(c.ctx, c.fallbackLogger, line, ch)
 		if err != nil {
 			return nil, err
@@ -300,7 +300,8 @@ func (c *rootCommand) setupLoggers() (<-chan struct{}, error) {
 		c.logFmt = "raw"
 
 	case strings.HasPrefix(line, "file"):
-		hook, err := log.FileHookFromConfigLine(c.ctx, c.fallbackLogger, line)
+		ch = make(chan struct{}) // TODO: refactor, get it from the constructor
+		hook, err := log.FileHookFromConfigLine(c.ctx, c.fallbackLogger, line, ch)
 		if err != nil {
 			return nil, err
 		}

diff --git a/log/file.go b/log/file.go
@@ -44,15 +44,19 @@ type fileHook struct {
 	w              io.WriteCloser
 	bw             *bufio.Writer
 	levels         []logrus.Level
+	done           chan struct{}
 }
 
 // FileHookFromConfigLine returns new fileHook hook.
 func FileHookFromConfigLine(
-	ctx context.Context, fallbackLogger logrus.FieldLogger, line string,
+	ctx context.Context, fallbackLogger logrus.FieldLogger, line string, done chan struct{},
 ) (logrus.Hook, error) {
+	// TODO: fix this so it works correctly with relative paths from the CWD
+
 	hook := &fileHook{
 		fallbackLogger: fallbackLogger,
 		levels:         logrus.AllLevels,
+		done:           done,
 	}
 
 	parts := strings.SplitN(line, "=", 2)
@@ -120,6 +124,7 @@ func (h *fileHook) loop(ctx context.Context) chan []byte {
 	loglines := make(chan []byte, fileHookBufferSize)
 
 	go func() {
+		defer close(h.done)
 		for {
 			select {
 			case entry := <-loglines:

diff --git a/log/file_test.go b/log/file_test.go
@@ -115,7 +115,9 @@ func TestFileHookFromConfigLine(t *testing.T) {
 		t.Run(test.line, func(t *testing.T) {
 			t.Parallel()
 
-			res, err := FileHookFromConfigLine(context.Background(), logrus.New(), test.line)
+			res, err := FileHookFromConfigLine(
+				context.Background(), logrus.New(), test.line, make(chan struct{}),
+			)
 
 			if test.err {
 				require.Error(t, err)
@@ -147,6 +149,7 @@ func TestFileHookFire(t *testing.T) {
 		w:        nc,
 		bw:       bufio.NewWriter(nc),
 		levels:   logrus.AllLevels,
+		done:     make(chan struct{}),
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())