add --max, --before, --after switches to dtail and dgrep commands

10 files changed, 77 insertions, 19 deletions

diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go
index 4da1bb3..4a6790b 100644
--- a/cmd/dgrep/main.go
+++ b/cmd/dgrep/main.go
@@ -36,12 +36,17 @@ func main() {
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&args.Discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&args.PrivateKeyPathFile, "key", "", "Path to private key")
-	flag.StringVar(&args.RegexStr, "regex", ".", "Regular expression")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&args.UserName, "user", userName, "Your system user name")
 	flag.StringVar(&args.What, "files", "", "File(s) to read")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
+
+	// Context awareness.
+	flag.StringVar(&args.LineContext.RegexStr, "regex", ".", "Regular expression")
 	flag.StringVar(&grep, "grep", "", "Alias for -regex")
+	flag.IntVar(&args.LineContext.BeforeContext, "before", 0, "Print lines of leading context before matching lines")
+	flag.IntVar(&args.LineContext.AfterContext, "after", 0, "Print lines of trailing context after matching lines")
+	flag.IntVar(&args.LineContext.MaxCount, "max", 0, "Stop reading file after NUM matching lines")
 
 	flag.Parse()
 
diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index f2a039f..1930483 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -50,14 +50,19 @@ func main() {
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&args.Discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&args.PrivateKeyPathFile, "key", "", "Path to private key")
-	flag.StringVar(&args.RegexStr, "regex", ".", "Regular expression")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&args.UserName, "user", userName, "Your system user name")
 	flag.StringVar(&args.What, "files", "", "File(s) to read")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
-	flag.StringVar(&grep, "grep", "", "Alias for -regex")
 	flag.StringVar(&queryStr, "query", "", "Map reduce query")
 
+	// Context awareness.
+	flag.StringVar(&args.LineContext.RegexStr, "regex", ".", "Regular expression")
+	flag.StringVar(&grep, "grep", "", "Alias for -regex")
+	flag.IntVar(&args.LineContext.BeforeContext, "before", 0, "Print lines of leading context before matching lines")
+	flag.IntVar(&args.LineContext.AfterContext, "after", 0, "Print lines of trailing context after matching lines")
+	flag.IntVar(&args.LineContext.MaxCount, "max", 0, "Stop reading file after NUM matching lines")
+
 	flag.Parse()
 
 	if grep != "" {
diff --git a/internal/clients/args.go b/internal/clients/args.go
index 7f782f1..67d2044 100644
--- a/internal/clients/args.go
+++ b/internal/clients/args.go
@@ -6,14 +6,22 @@ import (
 	gossh "golang.org/x/crypto/ssh"
 )
 
+// LineContext is here to help filtering out only specific lines.
+type LineContext struct {
+	RegexStr      string
+	AfterContext  int
+	BeforeContext int
+	MaxCount      int
+}
+
 // Args is a helper struct to summarize common client arguments.
 type Args struct {
+	LineContext
 	Mode               omode.Mode
 	ServersStr         string
 	UserName           string
 	What               string
 	Arguments          []string
-	RegexStr           string
 	RegexInvert        bool
 	TrustAllHosts      bool
 	Discovery          string
@@ -22,5 +30,5 @@ type Args struct {
 	SSHAuthMethods     []gossh.AuthMethod
 	SSHHostKeyCallback gossh.HostKeyCallback
 	PrivateKeyPathFile string
-	Quiet            bool
+	Quiet              bool
 }
diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index f20156f..0bdd62e 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -2,6 +2,8 @@ package clients
 
 import (
 	"context"
+	"fmt"
+	"strings"
 	"sync"
 	"time"
 
@@ -118,10 +120,44 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con
 	}
 }
 
+func (c *baseClient) makeCommandOptions() map[string]string {
+	options := make(map[string]string)
+
+	if c.Args.Quiet {
+		options["quiet"] = fmt.Sprintf("%v", c.Args.Quiet)
+	}
+	if c.Args.LineContext.MaxCount != 0 {
+		options["max"] = fmt.Sprintf("%d", c.Args.LineContext.MaxCount)
+	}
+	if c.Args.LineContext.BeforeContext != 0 {
+		options["before"] = fmt.Sprintf("%d", c.Args.LineContext.BeforeContext)
+	}
+	if c.Args.LineContext.AfterContext != 0 {
+		options["after"] = fmt.Sprintf("%d", c.Args.LineContext.AfterContext)
+	}
+
+	return options
+}
+
+func (c *baseClient) commandOptionsToString(options map[string]string) string {
+	var sb strings.Builder
+
+	count := 0
+	for k, v := range options {
+		if count > 0 {
+			sb.WriteString(":")
+		}
+		sb.WriteString(fmt.Sprintf("%s=%s", k, v))
+		count++
+	}
+
+	return sb.String()
+}
+
 func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) *remote.Connection {
 	conn := remote.NewConnection(server, c.UserName, sshAuthMethods, hostKeyCallback)
 	conn.Handler = c.maker.makeHandler(server)
-	conn.Commands = c.maker.makeCommands()
+	conn.Commands = c.maker.makeCommands(c.makeCommandOptions())
 
 	return conn
 }
diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go
index b7b6131..db892f1 100644
--- a/internal/clients/catclient.go
+++ b/internal/clients/catclient.go
@@ -41,10 +41,10 @@ func (c CatClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c CatClient) makeCommands() (commands []string) {
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
+func (c CatClient) makeCommands(options map[string]string) (commands []string) {
+	optionsStr := c.commandOptionsToString(options)
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
 	}
 	return
 }
diff --git a/internal/clients/grepclient.go b/internal/clients/grepclient.go
index 652c31b..567193a 100644
--- a/internal/clients/grepclient.go
+++ b/internal/clients/grepclient.go
@@ -40,10 +40,10 @@ func (c GrepClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c GrepClient) makeCommands() (commands []string) {
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
+func (c GrepClient) makeCommands(options map[string]string) (commands []string) {
+	optionsStr := c.commandOptionsToString(options)
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
 	}
 
 	return
diff --git a/internal/clients/maker.go b/internal/clients/maker.go
index d5ffd8b..a1d6864 100644
--- a/internal/clients/maker.go
+++ b/internal/clients/maker.go
@@ -9,5 +9,5 @@ import (
 // and send different commands to the DTail server.
 type maker interface {
 	makeHandler(server string) handlers.Handler
-	makeCommands() (commands []string)
+	makeCommands(options map[string]string) (commands []string)
 }
diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go
index 1c0c2cc..6a06aaf 100644
--- a/internal/clients/maprclient.go
+++ b/internal/clients/maprclient.go
@@ -110,21 +110,21 @@ func (c MaprClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewMaprHandler(server, c.query, c.globalGroup)
 }
 
-func (c MaprClient) makeCommands() (commands []string) {
+func (c MaprClient) makeCommands(options map[string]string) (commands []string) {
 	commands = append(commands, fmt.Sprintf("map %s", c.query.RawQuery))
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
 
 	modeStr := "cat"
 	if c.Mode == omode.TailClient {
 		modeStr = "tail"
 	}
+	optionsStr := c.commandOptionsToString(options)
 
 	for _, file := range strings.Split(c.What, ",") {
 		if c.Timeout > 0 {
 			commands = append(commands, fmt.Sprintf("timeout %d %s %s %s", c.Timeout, modeStr, file, c.Regex.Serialize()))
 			continue
 		}
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", modeStr, options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", modeStr, optionsStr, file, c.Regex.Serialize()))
 	}
 
 	return
diff --git a/internal/clients/tailclient.go b/internal/clients/tailclient.go
index cefbaa7..853ef1d 100644
--- a/internal/clients/tailclient.go
+++ b/internal/clients/tailclient.go
@@ -37,10 +37,10 @@ func (c TailClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c TailClient) makeCommands() (commands []string) {
-	options := fmt.Sprintf("quiet=%v", c.Args.Quiet)
+func (c TailClient) makeCommands(options map[string]string) (commands []string) {
+	optionsStr := c.commandOptionsToString(options)
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), options, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
 	}
 	logger.Debug(commands)
 
diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index 185e7c2..169c1eb 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -390,6 +390,10 @@ func (h *ServerHandler) decrementActiveReaders() int32 {
 	return atomic.LoadInt32(&h.activeReaders)
 }
 
+// TODO: All options related code should be in its own package (client + server)
+// Maybe we could move internal.clients.Args to internal.options.Options and
+// Use struct tagging to determine which ones should be serialized over the wire
+// from the client to the server.
 func readOptions(opts []string) (map[string]string, error) {
 	options := make(map[string]string, len(opts))