add dgrep context integration tests

7 files changed, 88 insertions, 10 deletions

diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go
index 2d7e53b..3cbb3cc 100644
--- a/cmd/dgrep/main.go
+++ b/cmd/dgrep/main.go
@@ -3,9 +3,14 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
 	"os"
 	"sync"
 
+	"net/http"
+	_ "net/http"
+	_ "net/http/pprof"
+
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
@@ -20,6 +25,7 @@ func main() {
 	var args config.Args
 	var displayVersion bool
 	var grep string
+	var pprof int
 	userName := user.Name()
 
 	flag.BoolVar(&args.NoColor, "noColor", false, "Disable ANSII terminal colors")
@@ -34,6 +40,7 @@ func main() {
 	flag.IntVar(&args.LContext.BeforeContext, "before", 0, "Print lines of leading context before matching lines")
 	flag.IntVar(&args.LContext.MaxCount, "max", 0, "Stop reading file after NUM matching lines")
 	flag.IntVar(&args.SSHPort, "port", config.DefaultSSHPort, "SSH server port")
+	flag.IntVar(&pprof, "pprof", -1, "Start PProf server this port")
 	flag.StringVar(&args.ConfigFile, "cfg", "", "Config file path")
 	flag.StringVar(&args.Discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&args.LogDir, "logDir", "~/log", "Log dir")
@@ -65,6 +72,13 @@ func main() {
 		args.RegexStr = grep
 	}
 
+	if pprof > -1 {
+		// For debugging purposes only
+		pprofArgs := fmt.Sprintf("0.0.0.0:%d", pprof)
+		go http.ListenAndServe(pprofArgs, nil)
+		dlog.Client.Info("Started PProf", pprofArgs)
+	}
+
 	client, err := clients.NewGrepClient(args)
 	if err != nil {
 		panic(err)
diff --git a/integrationtests/dgrep_test.go b/integrationtests/dgrep_test.go
index 57b5d86..b63ac45 100644
--- a/integrationtests/dgrep_test.go
+++ b/integrationtests/dgrep_test.go
@@ -47,3 +47,46 @@ func TestDGrep2(t *testing.T) {
 
 	os.Remove(stdoutFile)
 }
+
+func TestDGrepContext(t *testing.T) {
+	inFile := "mapr_testdata.log"
+	stdoutFile := "dgrepcontext.stdout.tmp"
+	expectedStdoutFile := "dgrepcontext.txt.expected"
+
+	_, err := runCommand(context.TODO(), t, stdoutFile,
+		"../dgrep", "--spartan", "--grep", "20211002-071947",
+		"-after", "3", "-before", "3", inFile)
+
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	if err := compareFiles(t, stdoutFile, expectedStdoutFile); err != nil {
+		t.Error(err)
+		return
+	}
+
+	os.Remove(stdoutFile)
+}
+
+func TestDGrepContext2(t *testing.T) {
+	inFile := "mapr_testdata.log"
+	stdoutFile := "dgrepcontext2.stdout.tmp"
+	expectedStdoutFile := "dgrepcontext2.txt.expected"
+
+	_, err := runCommand(context.TODO(), t, stdoutFile,
+		"../dgrep", "--spartan", "--grep", "20211002", "-max", "3", inFile)
+
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	if err := compareFiles(t, stdoutFile, expectedStdoutFile); err != nil {
+		t.Error(err)
+		return
+	}
+
+	os.Remove(stdoutFile)
+}
diff --git a/integrationtests/dgrepcontext.txt.expected b/integrationtests/dgrepcontext.txt.expected
new file mode 100644
index 0000000..ad0ae1f
--- /dev/null
+++ b/integrationtests/dgrepcontext.txt.expected
@@ -0,0 +1,9 @@
+INFO|20211002-071939|1|stats.go:56|8|11|7|0.80|471h8m17s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071946|1|stats.go:56|8|11|7|0.67|471h8m24s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071946|1|stats.go:56|8|11|7|0.67|471h8m24s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071947|1|stats.go:56|8|11|7|0.67|471h8m24s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071947|1|stats.go:56|8|11|7|0.67|471h8m25s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071947|1|stats.go:56|8|11|7|0.67|471h8m25s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071948|1|stats.go:56|8|11|7|0.67|471h8m25s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071948|1|stats.go:56|8|11|7|0.67|471h8m26s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=5
+INFO|20211002-071948|1|stats.go:56|8|15|7|0.67|471h8m26s|MAPREDUCE:STATS|currentConnections=1|lifetimeConnections=6
diff --git a/integrationtests/dgrepcontext2.txt.expected b/integrationtests/dgrepcontext2.txt.expected
new file mode 100644
index 0000000..38d4feb
--- /dev/null
+++ b/integrationtests/dgrepcontext2.txt.expected
@@ -0,0 +1,3 @@
+INFO|20211002-071143|1|stats.go:56|8|13|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1
+INFO|20211002-071143|1|stats.go:56|8|13|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1
+INFO|20211002-071143|1|stats.go:56|8|13|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1
diff --git a/integrationtests/dtail_test.go b/integrationtests/dtail_test.go
index 8e932a1..f180b8a 100644
--- a/integrationtests/dtail_test.go
+++ b/integrationtests/dtail_test.go
@@ -30,7 +30,7 @@ func TestDTailWithServer(t *testing.T) {
 	serverCh, _, _, err := startCommand(ctx, t,
 		"../dserver",
 		"--logger", "stdout",
-		"--logLevel", "info",
+		"--logLevel", "trace",
 		"--port", "4242",
 		"--relaxedAuth",
 	)
@@ -42,7 +42,7 @@ func TestDTailWithServer(t *testing.T) {
 	clientCh, _, _, err := startCommand(ctx, t,
 		"../dtail",
 		"--logger", "stdout",
-		"--logLevel", "devel",
+		"--logLevel", "trace",
 		"--servers", "localhost:4242",
 		"--files", followFile,
 		"--grep", "Hello",
diff --git a/internal/io/fs/readfile.go b/internal/io/fs/readfile.go
index 88d467e..28cbe58 100644
--- a/internal/io/fs/readfile.go
+++ b/internal/io/fs/readfile.go
@@ -99,15 +99,24 @@ func (f readFile) Start(ctx context.Context, ltx lcontext.LContext,
 	rawLines := make(chan *bytes.Buffer, 100)
 	truncate := make(chan struct{})
 
-	var wg sync.WaitGroup
-	wg.Add(1)
+	readCtx, readCancel := context.WithCancel(ctx)
+	var filterWg sync.WaitGroup
+	filterWg.Add(1)
 
 	go f.periodicTruncateCheck(ctx, truncate)
-	go f.filter(ctx, ltx, &wg, rawLines, lines, re)
+	go func() {
+		f.filter(ctx, ltx, rawLines, lines, re)
+		filterWg.Done()
+		// If the filter stopped, make the reader stop too, no need to read
+		// more data if there is nothing more the filter wants to filter for!
+		// E.g. it could be that we only want to filter N matches but not more.
+		readCancel()
+	}()
 
-	err = f.read(ctx, fd, rawLines, truncate)
+	err = f.read(readCtx, fd, rawLines, truncate)
 	close(rawLines)
-	wg.Wait()
+	// Filter may sends some data still. So wait until it is done here.
+	filterWg.Wait()
 
 	return err
 }
@@ -215,10 +224,8 @@ func (f readFile) read(ctx context.Context, fd *os.File, rawLines chan *bytes.Bu
 
 // Filter log lines matching a given regular expression.
 func (f readFile) filter(ctx context.Context, ltx lcontext.LContext,
-	wg *sync.WaitGroup, rawLines <-chan *bytes.Buffer, lines chan<- line.Line,
-	re regex.Regex) {
+	rawLines <-chan *bytes.Buffer, lines chan<- line.Line, re regex.Regex) {
 
-	defer wg.Done()
 	// Do we have any kind of local context settings? If so then run the more complex
 	// filterWithLContext method.
 	if ltx.Has() {
diff --git a/internal/server/handlers/basehandler.go b/internal/server/handlers/basehandler.go
index c25f85a..934f2bc 100644
--- a/internal/server/handlers/basehandler.go
+++ b/internal/server/handlers/basehandler.go
@@ -159,6 +159,8 @@ func (h *baseHandler) handleCommand(commandStr string) {
 		cancel()
 	}()
 
+	dlog.Server.Trace(args)
+	dlog.Server.Trace(args[0])
 	splitted := strings.Split(args[0], ":")
 	commandName := splitted[0]
 	options, ltx, err := config.DeserializeOptions(splitted[1:])