From d402cf103b00b8a55c6cab665ab36d4fd282f3a5 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Fri, 11 Sep 2020 11:34:12 +0100 Subject: no SIGINFO for now as not supported on Linux --- internal/io/signal/signal.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/io/signal/signal.go b/internal/io/signal/signal.go index bca7e6e..0d49485 100644 --- a/internal/io/signal/signal.go +++ b/internal/io/signal/signal.go @@ -10,7 +10,7 @@ import ( // StatsCh returns a channel for "please print stats" signalling. func StatsCh(ctx context.Context) <-chan struct{} { sigCh := make(chan os.Signal) - gosignal.Notify(sigCh, syscall.SIGINFO, syscall.SIGUSR1) + gosignal.Notify(sigCh, syscall.SIGUSR1) statsCh := make(chan struct{}) -- cgit v1.2.3 From ec67d9833095dfbe620dd3c99ea0caba391c4b87 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Fri, 11 Sep 2020 12:30:29 +0100 Subject: hit ctrl+c twice to exit --- cmd/dcat/main.go | 2 +- cmd/dgrep/main.go | 2 +- cmd/dmap/main.go | 2 +- cmd/drun/main.go | 2 +- cmd/dtail/main.go | 2 +- internal/clients/stats.go | 17 ++++++++++------- internal/io/signal/signal.go | 14 +++++++++++--- 7 files changed, 26 insertions(+), 15 deletions(-) diff --git a/cmd/dcat/main.go b/cmd/dcat/main.go index f0ea946..05e46ab 100644 --- a/cmd/dcat/main.go +++ b/cmd/dcat/main.go @@ -55,7 +55,7 @@ func main() { panic(err) } - status := client.Start(ctx, signal.StatsCh(ctx)) + status := client.Start(ctx, signal.InterruptCh(ctx)) logger.Flush() os.Exit(status) } diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go index d1fdc21..133631f 100644 --- a/cmd/dgrep/main.go +++ b/cmd/dgrep/main.go @@ -63,7 +63,7 @@ func main() { panic(err) } - status := client.Start(ctx, signal.StatsCh(ctx)) + status := client.Start(ctx, signal.InterruptCh(ctx)) logger.Flush() os.Exit(status) } diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go index 279b343..9f9ca9d 100644 --- a/cmd/dmap/main.go +++ b/cmd/dmap/main.go @@ -62,7 +62,7 @@ func main() { panic(err) } - status := client.Start(ctx, signal.StatsCh(ctx)) + status := client.Start(ctx, signal.InterruptCh(ctx)) logger.Flush() os.Exit(status) } diff --git a/cmd/drun/main.go b/cmd/drun/main.go index ffdf7bf..3916817 100644 --- a/cmd/drun/main.go +++ b/cmd/drun/main.go @@ -64,7 +64,7 @@ func main() { panic(err) } - status := client.Start(ctx, signal.StatsCh(ctx)) + status := client.Start(ctx, signal.InterruptCh(ctx)) logger.Flush() os.Exit(status) } diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go index ff9028b..aefaa6a 100644 --- a/cmd/dtail/main.go +++ b/cmd/dtail/main.go @@ -106,7 +106,7 @@ func main() { } } - status := client.Start(ctx, signal.StatsCh(ctx)) + status := client.Start(ctx, signal.InterruptCh(ctx)) logger.Flush() os.Exit(status) } diff --git a/internal/clients/stats.go b/internal/clients/stats.go index a6ac0c5..e7eabd8 100644 --- a/internal/clients/stats.go +++ b/internal/clients/stats.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "runtime" + "strings" "sync" "time" @@ -54,7 +55,8 @@ func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{}) if connected == connectedLast && !force { continue } - s.log(connected, newConnections, throttle) + + logger.Info(s.statsLine(connected, newConnections, throttle)) connectedLast = connected s.mutex.Lock() @@ -63,15 +65,16 @@ func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{}) } } -func (s *stats) log(connected, newConnections int, throttle int) { +func (s *stats) statsLine(connected, newConnections int, throttle int) string { percConnected := percentOf(float64(s.connectionsTotal), float64(connected)) - connectedStr := fmt.Sprintf("connected=%d/%d(%d%%)", connected, s.connectionsTotal, int(percConnected)) - newConnStr := fmt.Sprintf("new=%d", newConnections) - throttleStr := fmt.Sprintf("throttle=%d", throttle) - cpusGoroutinesStr := fmt.Sprintf("cpus/goroutines=%d/%d", runtime.NumCPU(), runtime.NumGoroutine()) + var stats []string + stats = append(stats, fmt.Sprintf("connected=%d/%d(%d%%)", connected, s.connectionsTotal, int(percConnected))) + stats = append(stats, fmt.Sprintf("new=%d", newConnections)) + stats = append(stats, fmt.Sprintf("throttle=%d", throttle)) + stats = append(stats, fmt.Sprintf("cpus/goroutines=%d/%d", runtime.NumCPU(), runtime.NumGoroutine())) - logger.Info("stats", connectedStr, newConnStr, throttleStr, cpusGoroutinesStr) + return strings.Join(stats, "|") } func (s *stats) numConnected() int { diff --git a/internal/io/signal/signal.go b/internal/io/signal/signal.go index 0d49485..3785d71 100644 --- a/internal/io/signal/signal.go +++ b/internal/io/signal/signal.go @@ -4,13 +4,15 @@ import ( "context" "os" gosignal "os/signal" - "syscall" + "time" + + "github.com/mimecast/dtail/internal/io/logger" ) // StatsCh returns a channel for "please print stats" signalling. -func StatsCh(ctx context.Context) <-chan struct{} { +func InterruptCh(ctx context.Context) <-chan struct{} { sigCh := make(chan os.Signal) - gosignal.Notify(sigCh, syscall.SIGUSR1) + gosignal.Notify(sigCh, os.Interrupt) statsCh := make(chan struct{}) @@ -20,6 +22,12 @@ func StatsCh(ctx context.Context) <-chan struct{} { case <-sigCh: select { case statsCh <- struct{}{}: + logger.Info("Hit Ctrl+C twice to exit") + select { + case <-sigCh: + os.Exit(0) + case <-time.After(time.Second): + } default: // Stats currently already printed. } -- cgit v1.2.3 From cdbcef72a1b9f93d08455b5a77f7f7afa2b9f53d Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Fri, 11 Sep 2020 14:55:50 +0100 Subject: fix nil pointer panic --- internal/clients/maprclient.go | 2 +- internal/regex/regex.go | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go index 581db44..149129d 100644 --- a/internal/clients/maprclient.go +++ b/internal/clients/maprclient.go @@ -123,7 +123,7 @@ func (c MaprClient) makeCommands() (commands []string) { 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 regex %s", modeStr, file, c.Regex.Serialize())) + commands = append(commands, fmt.Sprintf("%s %s %s", modeStr, file, c.Regex.Serialize())) } return diff --git a/internal/regex/regex.go b/internal/regex/regex.go index 707cb48..f668c38 100644 --- a/internal/regex/regex.go +++ b/internal/regex/regex.go @@ -39,6 +39,10 @@ func New(regexStr string, flag Flag) (Regex, error) { } func new(regexStr string, flags []Flag) (Regex, error) { + if len(flags) == 0 { + flags = append(flags, Default) + } + r := Regex{ regexStr: regexStr, flags: flags, -- cgit v1.2.3 From 813d2d00ec581c801d64091c7774988b559c3e93 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 17:52:45 +0100 Subject: refactor to have no context.Context in client handler structs --- internal/clients/baseclient.go | 2 +- internal/clients/handlers/basehandler.go | 22 +++++++++++++----- internal/clients/handlers/clienthandler.go | 4 +--- internal/clients/handlers/done.go | 37 ++++++++++++++++++++++++++++++ internal/clients/handlers/handler.go | 3 +-- internal/clients/handlers/healthhandler.go | 15 ++++++++---- internal/clients/handlers/maprhandler.go | 4 +--- internal/clients/handlers/withcancel.go | 24 ------------------- internal/clients/healthclient.go | 2 +- internal/clients/remote/connection.go | 7 +++--- 10 files changed, 73 insertions(+), 47 deletions(-) create mode 100644 internal/clients/handlers/done.go delete mode 100644 internal/clients/handlers/withcancel.go diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index 008a01e..d8d4fde 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -99,7 +99,7 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con defer func() { <-active }() for { - connCtx, cancel := conn.Handler.WithCancel(ctx) + connCtx, cancel := context.WithCancel(ctx) defer cancel() conn.Start(connCtx, cancel, c.throttleCh, c.stats.connectionsEstCh) diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 65bbfd7..54b80ae 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -13,7 +13,7 @@ import ( ) type baseHandler struct { - withCancel + done *Done server string shellStarted bool commands chan string @@ -29,6 +29,14 @@ func (h *baseHandler) Status() int { return h.status } +func (h *baseHandler) Done() <-chan struct{} { + return h.done.Done() +} + +func (h *baseHandler) Shutdown() { + h.done.Shutdown() +} + // SendMessage to the server. func (h *baseHandler) SendMessage(command string) error { encoded := base64.StdEncoding.EncodeToString([]byte(command)) @@ -38,7 +46,8 @@ func (h *baseHandler) SendMessage(command string) error { case h.commands <- fmt.Sprintf("protocol %s base64 %v;", version.ProtocolCompat, encoded): case <-time.After(time.Second * 5): return fmt.Errorf("Timed out sending command '%s' (base64: '%s')", command, encoded) - case <-h.ctx.Done(): + case <-h.Done(): + return nil } return nil @@ -65,7 +74,7 @@ func (h *baseHandler) Read(p []byte) (n int, err error) { select { case command := <-h.commands: n = copy(p, []byte(command)) - case <-h.ctx.Done(): + case <-h.Done(): return 0, io.EOF } return @@ -95,10 +104,11 @@ func (h *baseHandler) handleHiddenMessage(message string) { case strings.HasPrefix(message, ".syn close connection"): h.SendMessage(".ack close connection") select { - case <-time.After(time.Second * 1): + case <-time.After(time.Second * 5): logger.Debug("Shutting down client after timeout and sending ack to server") - h.withCancel.shutdown() - case <-h.ctx.Done(): + h.Shutdown() + case <-h.Done(): + return } case strings.HasPrefix(message, ".run exitstatus"): diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go index fcd8052..2908f7c 100644 --- a/internal/clients/handlers/clienthandler.go +++ b/internal/clients/handlers/clienthandler.go @@ -19,9 +19,7 @@ func NewClientHandler(server string) *ClientHandler { shellStarted: false, commands: make(chan string), status: -1, - withCancel: withCancel{ - done: make(chan struct{}), - }, + done: NewDone(), }, } } diff --git a/internal/clients/handlers/done.go b/internal/clients/handlers/done.go new file mode 100644 index 0000000..5b1335e --- /dev/null +++ b/internal/clients/handlers/done.go @@ -0,0 +1,37 @@ +package handlers + +import ( + "sync" + + "github.com/mimecast/dtail/internal/io/logger" +) + +type Done struct { + ch chan struct{} + mutex sync.Mutex +} + +func NewDone() *Done { + return &Done{ + ch: make(chan struct{}), + } +} + +func (d *Done) Done() <-chan struct{} { + return d.ch +} + +func (d *Done) Shutdown() { + d.mutex.Lock() + defer d.mutex.Unlock() + + logger.Debug("Done.Shutdown()") + + select { + case <-d.ch: + return + default: + logger.Debug("Done.Shutdown() -> close") + close(d.ch) + } +} diff --git a/internal/clients/handlers/handler.go b/internal/clients/handlers/handler.go index c53ca34..afa87e2 100644 --- a/internal/clients/handlers/handler.go +++ b/internal/clients/handlers/handler.go @@ -1,7 +1,6 @@ package handlers import ( - "context" "io" ) @@ -11,6 +10,6 @@ type Handler interface { SendMessage(command string) error Server() string Status() int - WithCancel(ctx context.Context) (context.Context, context.CancelFunc) + Shutdown() Done() <-chan struct{} } diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go index 9051015..9fc2671 100644 --- a/internal/clients/handlers/healthhandler.go +++ b/internal/clients/handlers/healthhandler.go @@ -8,7 +8,7 @@ import ( // HealthHandler implements the handler required for health checks. type HealthHandler struct { - withCancel + done *Done // Buffer of incoming data from server. receiveBuf []byte // To send commands to the server. @@ -27,9 +27,7 @@ func NewHealthHandler(server string, receive chan<- string) *HealthHandler { receive: receive, commands: make(chan string), status: -1, - withCancel: withCancel{ - done: make(chan struct{}), - }, + done: NewDone(), } return &h @@ -45,12 +43,21 @@ func (h *HealthHandler) Status() int { return h.status } +func (h *HealthHandler) Done() <-chan struct{} { + return h.done.Done() +} + +func (h *HealthHandler) Shutdown() { + h.done.Shutdown() +} + // SendMessage sends a DTail command to the server. func (h *HealthHandler) SendMessage(command string) error { select { case h.commands <- fmt.Sprintf("%s;", command): case <-time.NewTimer(time.Second * 10).C: return errors.New("Timed out sending command " + command) + case <-h.Done(): } return nil diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go index b908f3b..5d98690 100644 --- a/internal/clients/handlers/maprhandler.go +++ b/internal/clients/handlers/maprhandler.go @@ -24,9 +24,7 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr shellStarted: false, commands: make(chan string), status: -1, - withCancel: withCancel{ - done: make(chan struct{}), - }, + done: NewDone(), }, query: query, aggregate: client.NewAggregate(server, query, globalGroup), diff --git a/internal/clients/handlers/withcancel.go b/internal/clients/handlers/withcancel.go deleted file mode 100644 index 7c9cf4e..0000000 --- a/internal/clients/handlers/withcancel.go +++ /dev/null @@ -1,24 +0,0 @@ -package handlers - -import "context" - -type withCancel struct { - ctx context.Context - done chan struct{} -} - -// WithCancel sets and returns the context used. -func (w *withCancel) WithCancel(ctx context.Context) (context.Context, context.CancelFunc) { - cancelCtx, cancel := context.WithCancel(ctx) - w.ctx = cancelCtx - - return cancelCtx, cancel -} - -func (w *withCancel) Done() <-chan struct{} { - return w.done -} - -func (w *withCancel) shutdown() { - close(w.done) -} diff --git a/internal/clients/healthclient.go b/internal/clients/healthclient.go index 7313583..e93f6be 100644 --- a/internal/clients/healthclient.go +++ b/internal/clients/healthclient.go @@ -50,7 +50,7 @@ func (c *HealthClient) Start(ctx context.Context) (status int) { conn.Handler = handlers.NewHealthHandler(c.server, receive) conn.Commands = []string{c.mode.String()} - connCtx, cancel := conn.Handler.WithCancel(ctx) + connCtx, cancel := context.WithCancel(ctx) go conn.Start(connCtx, cancel, throttleCh, statsCh) for { diff --git a/internal/clients/remote/connection.go b/internal/clients/remote/connection.go index 2d97d14..b29ffed 100644 --- a/internal/clients/remote/connection.go +++ b/internal/clients/remote/connection.go @@ -177,21 +177,21 @@ func (c *Connection) handle(ctx context.Context, cancel context.CancelFunc, sess } go func() { - defer cancel() io.Copy(stdinPipe, c.Handler) + cancel() }() go func() { - defer cancel() io.Copy(c.Handler, stdoutPipe) + cancel() }() go func() { - defer cancel() select { case <-c.Handler.Done(): case <-ctx.Done(): } + cancel() }() // Send all commands to client. @@ -207,5 +207,6 @@ func (c *Connection) handle(ctx context.Context, cancel context.CancelFunc, sess } <-ctx.Done() + c.Handler.Shutdown() return nil } -- cgit v1.2.3 From bad8e04bb4410b94b8e875ccde287f74ab94121a Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 18:38:36 +0100 Subject: server handler context refactoring --- internal/clients/handlers/basehandler.go | 3 +- internal/clients/handlers/clienthandler.go | 3 +- internal/clients/handlers/done.go | 37 --------------- internal/clients/handlers/healthhandler.go | 6 ++- internal/clients/handlers/maprhandler.go | 3 +- internal/done.go | 32 +++++++++++++ internal/server/handlers/controlhandler.go | 26 +++++++---- internal/server/handlers/handler.go | 2 + internal/server/handlers/serverhandler.go | 73 ++++++++++++++++-------------- internal/server/server.go | 39 +++++++--------- 10 files changed, 116 insertions(+), 108 deletions(-) delete mode 100644 internal/clients/handlers/done.go create mode 100644 internal/done.go diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go index 54b80ae..b5045e2 100644 --- a/internal/clients/handlers/basehandler.go +++ b/internal/clients/handlers/basehandler.go @@ -8,12 +8,13 @@ import ( "strings" "time" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/version" ) type baseHandler struct { - done *Done + done *internal.Done server string shellStarted bool commands chan string diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go index 2908f7c..2bcb038 100644 --- a/internal/clients/handlers/clienthandler.go +++ b/internal/clients/handlers/clienthandler.go @@ -1,6 +1,7 @@ package handlers import ( + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" ) @@ -19,7 +20,7 @@ func NewClientHandler(server string) *ClientHandler { shellStarted: false, commands: make(chan string), status: -1, - done: NewDone(), + done: internal.NewDone(), }, } } diff --git a/internal/clients/handlers/done.go b/internal/clients/handlers/done.go deleted file mode 100644 index 5b1335e..0000000 --- a/internal/clients/handlers/done.go +++ /dev/null @@ -1,37 +0,0 @@ -package handlers - -import ( - "sync" - - "github.com/mimecast/dtail/internal/io/logger" -) - -type Done struct { - ch chan struct{} - mutex sync.Mutex -} - -func NewDone() *Done { - return &Done{ - ch: make(chan struct{}), - } -} - -func (d *Done) Done() <-chan struct{} { - return d.ch -} - -func (d *Done) Shutdown() { - d.mutex.Lock() - defer d.mutex.Unlock() - - logger.Debug("Done.Shutdown()") - - select { - case <-d.ch: - return - default: - logger.Debug("Done.Shutdown() -> close") - close(d.ch) - } -} diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go index 9fc2671..95693ab 100644 --- a/internal/clients/handlers/healthhandler.go +++ b/internal/clients/handlers/healthhandler.go @@ -4,11 +4,13 @@ import ( "errors" "fmt" "time" + + "github.com/mimecast/dtail/internal" ) // HealthHandler implements the handler required for health checks. type HealthHandler struct { - done *Done + done *internal.Done // Buffer of incoming data from server. receiveBuf []byte // To send commands to the server. @@ -27,7 +29,7 @@ func NewHealthHandler(server string, receive chan<- string) *HealthHandler { receive: receive, commands: make(chan string), status: -1, - done: NewDone(), + done: internal.NewDone(), } return &h diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go index 5d98690..fb71c8f 100644 --- a/internal/clients/handlers/maprhandler.go +++ b/internal/clients/handlers/maprhandler.go @@ -3,6 +3,7 @@ package handlers import ( "strings" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/mapr" "github.com/mimecast/dtail/internal/mapr/client" @@ -24,7 +25,7 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr shellStarted: false, commands: make(chan string), status: -1, - done: NewDone(), + done: internal.NewDone(), }, query: query, aggregate: client.NewAggregate(server, query, globalGroup), diff --git a/internal/done.go b/internal/done.go new file mode 100644 index 0000000..2326eee --- /dev/null +++ b/internal/done.go @@ -0,0 +1,32 @@ +package internal + +import ( + "sync" +) + +type Done struct { + ch chan struct{} + mutex sync.Mutex +} + +func NewDone() *Done { + return &Done{ + ch: make(chan struct{}), + } +} + +func (d *Done) Done() <-chan struct{} { + return d.ch +} + +func (d *Done) Shutdown() { + d.mutex.Lock() + defer d.mutex.Unlock() + + select { + case <-d.ch: + return + default: + close(d.ch) + } +} diff --git a/internal/server/handlers/controlhandler.go b/internal/server/handlers/controlhandler.go index daa9835..9a8eb75 100644 --- a/internal/server/handlers/controlhandler.go +++ b/internal/server/handlers/controlhandler.go @@ -1,20 +1,19 @@ package handlers import ( - "context" "fmt" "io" "os" "strings" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/logger" user "github.com/mimecast/dtail/internal/user/server" ) // ControlHandler is used for control functions and health monitoring. type ControlHandler struct { - ctx context.Context - done chan struct{} + done *internal.Done hostname string payload []byte serverMessages chan string @@ -22,12 +21,11 @@ type ControlHandler struct { } // NewControlHandler returns a new control handler. -func NewControlHandler(ctx context.Context, user *user.User) (*ControlHandler, <-chan struct{}) { +func NewControlHandler(user *user.User) *ControlHandler { logger.Debug(user, "Creating control handler") h := ControlHandler{ - ctx: ctx, - done: make(chan struct{}), + done: internal.NewDone(), serverMessages: make(chan string, 10), user: user, } @@ -40,7 +38,15 @@ func NewControlHandler(ctx context.Context, user *user.User) (*ControlHandler, < s := strings.Split(fqdn, ".") h.hostname = s[0] - return &h, h.done + return &h +} + +func (h *ControlHandler) Shutdown() { + h.done.Shutdown() +} + +func (h *ControlHandler) Done() <-chan struct{} { + return h.done.Done() } // Read is to send data to the client via the Reader interface. @@ -51,7 +57,7 @@ func (h *ControlHandler) Read(p []byte) (n int, err error) { wholePayload := []byte(fmt.Sprintf("SERVER|%s|%s\n", h.hostname, message)) n = copy(p, wholePayload) return - case <-h.ctx.Done(): + case <-h.done.Done(): return 0, io.EOF } } @@ -63,7 +69,7 @@ func (h *ControlHandler) Write(p []byte) (n int, err error) { switch c { case ';': wholePayload := strings.TrimSpace(string(h.payload)) - h.handleCommand(h.ctx, wholePayload) + h.handleCommand(wholePayload) h.payload = nil default: @@ -75,7 +81,7 @@ func (h *ControlHandler) Write(p []byte) (n int, err error) { return } -func (h *ControlHandler) handleCommand(ctx context.Context, command string) { +func (h *ControlHandler) handleCommand(command string) { logger.Info(h.user, command) s := strings.Split(command, " ") logger.Debug(h.user, "Receiving command", command, s) diff --git a/internal/server/handlers/handler.go b/internal/server/handlers/handler.go index c42ceb9..b04e854 100644 --- a/internal/server/handlers/handler.go +++ b/internal/server/handlers/handler.go @@ -5,4 +5,6 @@ import "io" // Handler interface for server side functionality. type Handler interface { io.ReadWriter + Shutdown() + Done() <-chan struct{} } diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 7017f3e..5af3ebb 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -13,6 +13,7 @@ import ( "sync/atomic" "time" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/line" "github.com/mimecast/dtail/internal/io/logger" @@ -31,33 +32,28 @@ const ( // the Bi-directional communication between SSH client and server. // This handler implements the handler of the SSH server. type ServerHandler struct { - lines chan line.Line - regex string - aggregate *server.Aggregate - aggregatedMessages chan string - serverMessages chan string - payload []byte - hostname string - user *user.User - // TODO: Move all these channels into a separate struct for readability! + done *internal.Done + lines chan line.Line + regex string + aggregate *server.Aggregate + aggregatedMessages chan string + serverMessages chan string + payload []byte + hostname string + user *user.User catLimiter chan struct{} tailLimiter chan struct{} globalServerWaitFor chan struct{} ackCloseReceived chan struct{} - serverCtx context.Context - handlerCtx context.Context - done chan struct{} activeCommands int32 activeReaders int32 background background.Background } // NewServerHandler returns the server handler. -func NewServerHandler(handlerCtx, serverCtx context.Context, user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}, background background.Background) (*ServerHandler, <-chan struct{}) { +func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}, background background.Background) *ServerHandler { h := ServerHandler{ - serverCtx: serverCtx, - handlerCtx: handlerCtx, - done: make(chan struct{}), + done: internal.NewDone(), lines: make(chan line.Line, 100), serverMessages: make(chan string, 10), aggregatedMessages: make(chan string, 10), @@ -78,7 +74,15 @@ func NewServerHandler(handlerCtx, serverCtx context.Context, user *user.User, ca s := strings.Split(fqdn, ".") h.hostname = s[0] - return &h, h.done + return &h +} + +func (h *ServerHandler) Shutdown() { + h.done.Shutdown() +} + +func (h *ServerHandler) Done() <-chan struct{} { + return h.done.Done() } // Read is to send data to the dtail client via Reader interface. @@ -120,7 +124,7 @@ func (h *ServerHandler) Read(p []byte) (n int, err error) { case <-time.After(time.Second): // Once in a while check whether we are done. select { - case <-h.handlerCtx.Done(): + case <-h.done.Done(): return 0, io.EOF default: } @@ -134,7 +138,7 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) { switch c { case ';': commandStr := strings.TrimSpace(string(h.payload)) - h.handleCommand(h.handlerCtx, commandStr) + h.handleCommand(commandStr) h.payload = nil default: h.payload = append(h.payload, c) @@ -145,9 +149,10 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) { return } -func (h *ServerHandler) handleCommand(ctx context.Context, commandStr string) { +func (h *ServerHandler) handleCommand(commandStr string) { logger.Debug(h.user, commandStr) var timeout time.Duration + ctx := context.Background() args, argc, err := h.handleProtocolVersion(strings.Split(commandStr, " ")) if err != nil { @@ -172,15 +177,21 @@ func (h *ServerHandler) handleCommand(ctx context.Context, commandStr string) { return } + ctx, cancel := context.WithCancel(ctx) + go func() { + <-h.done.Done() + cancel() + }() + if timeout > 0 { logger.Info(h.user, "Command with timeout context", argc, args, timeout) - commandCtx, cancel := context.WithTimeout(ctx, timeout) + ctx, cancel := context.WithTimeout(ctx, timeout) go func() { - <-commandCtx.Done() + <-ctx.Done() logger.Info(h.user, "Command timed out, canceling it", args, args, timeout) cancel() }() - h.handleUserCommand(commandCtx, argc, args, timeout) + h.handleUserCommand(ctx, argc, args, timeout) return } @@ -348,9 +359,8 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args [] // Set default background timeout. timeout = time.Hour * 1 } - // Use a new context based on the server context, so that background job does not get - // terminated when handler/SSH connection terminates. - commandCtx, cancel := context.WithTimeout(h.serverCtx, timeout) + + commandCtx, cancel := context.WithTimeout(ctx, timeout) if err := h.background.Add(h.user.Name, jobName, cancel, &wg); err != nil { h.sendServerMessage(logger.Error(h.user, err, jobName, args)) @@ -406,7 +416,7 @@ func (h *ServerHandler) handleAckCommand(argc int, args []string) { func (h *ServerHandler) send(ch chan<- string, message string) { select { case ch <- message: - case <-h.handlerCtx.Done(): + case <-h.done.Done(): } } @@ -447,7 +457,7 @@ func (h *ServerHandler) shutdown() { go func() { select { case h.serverMessageC() <- ".syn close connection": - case <-h.handlerCtx.Done(): + case <-h.done.Done(): } }() @@ -455,13 +465,10 @@ func (h *ServerHandler) shutdown() { case <-h.ackCloseReceived: case <-time.After(time.Second * 5): logger.Debug(h.user, "Shutdown timeout reached, enforcing shutdown") - case <-h.handlerCtx.Done(): + case <-h.done.Done(): } - select { - case h.done <- struct{}{}: - default: - } + h.done.Shutdown() } func (h *ServerHandler) incrementActiveCommands() { diff --git a/internal/server/server.go b/internal/server/server.go index a446738..5e2a521 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -178,53 +178,46 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch switch req.Type { case "shell": - handlerCtx, cancel := context.WithCancel(ctx) - var handler handlers.Handler - var done <-chan struct{} - switch user.Name { case config.ControlUser: - handler, done = handlers.NewControlHandler(handlerCtx, user) + handler = handlers.NewControlHandler(user) default: - handler, done = handlers.NewServerHandler(handlerCtx, ctx, user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor, s.background) + handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor, s.background) } - go func() { - // Handler finished work, cancel all remaining routines - defer cancel() - - <-done - }() + terminate := func() { + handler.Shutdown() + sshConn.Close() + } go func() { // Broken pipe, cancel - defer cancel() - io.Copy(channel, handler) + terminate() }() go func() { // Broken pipe, cancel - defer cancel() - io.Copy(handler, channel) + terminate() }() go func() { - defer cancel() + select { + case <-ctx.Done(): + case <-handler.Done(): + } + terminate() + }() + go func() { if err := sshConn.Wait(); err != nil && err != io.EOF { logger.Error(user, err) } s.stats.decrementConnections() logger.Info(user, "Good bye Mister!") - }() - - go func() { - <-handlerCtx.Done() - sshConn.Close() - logger.Info(user, "Closed SSH connection") + terminate() }() // Only serving shell type -- cgit v1.2.3 From df2ff83897cde61d04b12958c6f6d458c69502f4 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 19:15:08 +0100 Subject: refactor to have no context.Context in server mapreduce aggregate structs --- internal/mapr/server/aggregate.go | 80 +++++++++++++++++-------------- internal/server/handlers/serverhandler.go | 2 +- 2 files changed, 44 insertions(+), 38 deletions(-) diff --git a/internal/mapr/server/aggregate.go b/internal/mapr/server/aggregate.go index 1028943..cd59b63 100644 --- a/internal/mapr/server/aggregate.go +++ b/internal/mapr/server/aggregate.go @@ -6,6 +6,7 @@ import ( "strings" "time" + "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/line" "github.com/mimecast/dtail/internal/io/logger" @@ -15,6 +16,7 @@ import ( // Aggregate is for aggregating mapreduce data on the DTail server side. type Aggregate struct { + done *internal.Done // Log lines to process (parsing MAPREDUCE lines). Lines chan line.Line // Hostname of the current server (used to populate $hostname field). @@ -23,12 +25,12 @@ type Aggregate struct { serialize chan struct{} // Signals to flush data. flush chan struct{} + // Signals that data has been flushed + flushed chan struct{} // The mapr query query *mapr.Query // The mapr log format parser parser *logformat.Parser - cancel context.CancelFunc - ctx context.Context } // NewAggregate return a new server side aggregator. @@ -64,56 +66,63 @@ func NewAggregate(queryStr string) (*Aggregate, error) { } } - ctx, cancel := context.WithCancel(context.Background()) - a := Aggregate{ + done: internal.NewDone(), Lines: make(chan line.Line, 100), serialize: make(chan struct{}), flush: make(chan struct{}), + flushed: make(chan struct{}), hostname: s[0], query: query, parser: logParser, - ctx: ctx, - cancel: cancel, } return &a, nil } +func (a *Aggregate) Shutdown() { + a.Flush() + a.done.Shutdown() +} + // Start an aggregation. func (a *Aggregate) Start(ctx context.Context, maprLines chan<- string) { - defer a.cancel() - fieldsCh := a.linesToFields(ctx) + myCtx, cancel := context.WithCancel(ctx) + defer cancel() + + go func() { + select { + case <-myCtx.Done(): + a.done.Shutdown() + case <-a.done.Done(): + cancel() + } + }() + + fieldsCh := a.makeFields(myCtx) // Add fields (e.g. via 'set' clause) if len(a.query.Set) > 0 { - fieldsCh = a.addMoreFields(ctx, fieldsCh) + fieldsCh = a.addFields(myCtx, fieldsCh) } - go a.fieldsToMaprLines(ctx, fieldsCh, maprLines) - a.periodicAggregateTimer(ctx) + go a.aggregateTimer(myCtx) + a.makeMaprLines(myCtx, fieldsCh, maprLines) } -// Cancel the aggregation. -func (a *Aggregate) Cancel() { - a.cancel() -} - -func (a *Aggregate) periodicAggregateTimer(ctx context.Context) { +func (a *Aggregate) aggregateTimer(ctx context.Context) { for { select { case <-time.After(a.query.Interval): a.Serialize(ctx) case <-ctx.Done(): return - case <-a.ctx.Done(): - return } } } -func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string { +func (a *Aggregate) makeFields(ctx context.Context) <-chan map[string]string { ch := make(chan map[string]string) go func() { @@ -144,8 +153,6 @@ func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string } case <-ctx.Done(): return - case <-a.ctx.Done(): - return } } }() @@ -153,14 +160,14 @@ func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string return ch } -func (a *Aggregate) addMoreFields(ctx context.Context, fieldsCh <-chan map[string]string) <-chan map[string]string { +func (a *Aggregate) addFields(ctx context.Context, fieldsCh <-chan map[string]string) <-chan map[string]string { ch := make(chan map[string]string) go func() { defer close(ch) for { - // fieldsCh will be closed via 'linesToFields' if ctx is done + // fieldsCh will be closed via 'makeFields' if ctx is done fields, ok := <-fieldsCh if !ok { return @@ -179,7 +186,7 @@ func (a *Aggregate) addMoreFields(ctx context.Context, fieldsCh <-chan map[strin return ch } -func (a *Aggregate) fieldsToMaprLines(ctx context.Context, fieldsCh <-chan map[string]string, maprLines chan<- string) { +func (a *Aggregate) makeMaprLines(ctx context.Context, fieldsCh <-chan map[string]string, maprLines chan<- string) { group := mapr.NewGroupSet() serialize := func() { @@ -200,18 +207,10 @@ func (a *Aggregate) fieldsToMaprLines(ctx context.Context, fieldsCh <-chan map[s case <-a.serialize: serialize() case <-a.flush: - logger.Info("Flushing mapreduce result") serialize() - a.flush <- struct{}{} - logger.Info("Done flushing mapreduce result") + a.flushed <- struct{}{} case <-ctx.Done(): return - case <-a.ctx.Done(): - logger.Info("Flushing mapreduce result") - serialize() - a.flush <- struct{}{} - logger.Info("Done flushing mapreduce result") - return } } } @@ -254,6 +253,8 @@ func (a *Aggregate) aggregate(group *mapr.GroupSet, fields map[string]string) { func (a *Aggregate) Serialize(ctx context.Context) { select { case a.serialize <- struct{}{}: + case <-time.After(time.Minute): + logger.Warn("Starting to serialize mapredice data takes over a minute") case <-ctx.Done(): } } @@ -261,15 +262,20 @@ func (a *Aggregate) Serialize(ctx context.Context) { // Flush all data. func (a *Aggregate) Flush() { select { - case <-a.ctx.Done(): - return case a.flush <- struct{}{}: + logger.Info("Flushing mapreduce data") case <-time.After(time.Minute): + logger.Warn("Starting to flush mapreduce data takes over a minute") + return + case <-a.done.Done(): return } select { - case <-a.flush: + case <-a.flushed: + logger.Info("Done flushing") case <-time.After(time.Minute): + logger.Warn("Waiting for data to be flushed takes over a minute") + case <-a.done.Done(): } } diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 5af3ebb..164a280 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -266,7 +266,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args [] if h.aggregate == nil { return } - h.aggregate.Cancel() + h.aggregate.Shutdown() } } -- cgit v1.2.3 From 98ed1491676d1e69dfbd6126ce1761a473951e9f Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sat, 19 Sep 2020 20:16:45 +0100 Subject: INFO|INFO -> INFO --- internal/io/logger/logger.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/io/logger/logger.go b/internal/io/logger/logger.go index d059cbb..b7db0a7 100644 --- a/internal/io/logger/logger.go +++ b/internal/io/logger/logger.go @@ -224,7 +224,7 @@ func log(what string, severity string, args []interface{}) string { return "" } - messages := []string{severity} + messages := []string{} for _, arg := range args { switch v := arg.(type) { -- cgit v1.2.3 From 7df612f527bd5dc2e785bf766d7d61124c260b94 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 16 Nov 2020 22:11:44 +0000 Subject: remove drun command for simplicity. only focus on interactive commands dealing with log streams --- Makefile | 2 - cmd/drun/main.go | 88 ------------------ doc/drun.gif | Bin 1668728 -> 0 bytes doc/examples.md | 11 --- doc/quickstart.md | 3 +- internal/clients/runclient.go | 87 ----------------- internal/io/run/run.go | 150 ------------------------------ internal/server/background/background.go | 126 ------------------------- internal/server/handlers/runcommand.go | 111 ---------------------- internal/server/handlers/serverhandler.go | 85 +---------------- internal/server/server.go | 6 +- 11 files changed, 3 insertions(+), 666 deletions(-) delete mode 100644 cmd/drun/main.go delete mode 100644 doc/drun.gif delete mode 100644 internal/clients/runclient.go delete mode 100644 internal/io/run/run.go delete mode 100644 internal/server/background/background.go delete mode 100644 internal/server/handlers/runcommand.go diff --git a/Makefile b/Makefile index 75b9333..7cf125f 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,6 @@ build: ${GO} build -o dcat ./cmd/dcat/main.go ${GO} build -o dgrep ./cmd/dgrep/main.go ${GO} build -o dmap ./cmd/dmap/main.go - ${GO} build -o drun ./cmd/drun/main.go ${GO} build -o dtail ./cmd/dtail/main.go clean: ls ./cmd/ | while read cmd; do \ @@ -16,7 +15,6 @@ install: build cp -pv dcat ${GOPATH}/bin/dcat cp -pv dgrep ${GOPATH}/bin/dgrep cp -pv dmap ${GOPATH}/bin/dmap - cp -pv drun ${GOPATH}/bin/drun cp -pv dtail ${GOPATH}/bin/dtail vet: find . -type d | while read dir; do \ diff --git a/cmd/drun/main.go b/cmd/drun/main.go deleted file mode 100644 index 3916817..0000000 --- a/cmd/drun/main.go +++ /dev/null @@ -1,88 +0,0 @@ -package main - -import ( - "context" - "flag" - "io/ioutil" - "os" - "strings" - - "github.com/mimecast/dtail/internal/clients" - "github.com/mimecast/dtail/internal/color" - "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/io/signal" - "github.com/mimecast/dtail/internal/user" - "github.com/mimecast/dtail/internal/version" -) - -// The evil begins here. -func main() { - var args clients.Args - var background string - var cfgFile string - var command string - var debugEnable bool - var displayVersion bool - var jobName string - var noColor bool - var sshPort int - - userName := user.Name() - - flag.BoolVar(&args.TrustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys") - flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages") - flag.BoolVar(&displayVersion, "version", false, "Display version") - flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors") - flag.IntVar(&args.ConnectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently") - flag.IntVar(&args.Timeout, "timeout", 0, "Command execution timeout") - 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.ServersStr, "servers", "", "Remote servers to connect") - flag.StringVar(&args.UserName, "user", userName, "Your system user name") - flag.StringVar(&background, "background", "", "Can be one of 'start', 'cancel', 'list' or empty") - flag.StringVar(&cfgFile, "cfg", "", "Config file path") - flag.StringVar(&command, "command", "", "Command to run") - flag.StringVar(&jobName, "name", "", "The job name (if run in background)") - - flag.Parse() - - config.Read(cfgFile, sshPort) - color.Colored = !noColor - - if displayVersion { - version.PrintAndExit() - } - - ctx := context.TODO() - logger.Start(ctx, logger.Modes{Debug: debugEnable || config.Common.DebugEnable}) - - args.What, args.Arguments = readCommand(command) - client, err := clients.NewRunClient(args, background, jobName) - if err != nil { - panic(err) - } - - status := client.Start(ctx, signal.InterruptCh(ctx)) - logger.Flush() - os.Exit(status) -} - -func readCommand(command string) (string, []string) { - splitted := strings.Split(command, " ") - - script := splitted[0] - if _, err := os.Stat(script); os.IsNotExist(err) { - var commandArgs []string - return command, commandArgs - } - commandArgs := splitted[1:] - - bytes, err := ioutil.ReadFile(script) - if err != nil { - panic(err) - } - - return string(bytes), commandArgs -} diff --git a/doc/drun.gif b/doc/drun.gif deleted file mode 100644 index a3efeed..0000000 Binary files a/doc/drun.gif and /dev/null differ diff --git a/doc/examples.md b/doc/examples.md index 5a5d892..964660a 100644 --- a/doc/examples.md +++ b/doc/examples.md @@ -52,17 +52,6 @@ The following example demonstrates how to grep files (display only the lines whi ![dgrep](dgrep.gif "Grep example") -# How to use ``drun`` - -The following example demonstrates how to execute a command on multiple machines remotely: - -```shell -% drun --servers <(head -n 30 serverlist.txt) \ - --command uptime -``` - -![dgrep](drun.gif "Run example") - # How to use ``dmap`` To run a mapreduce aggregation over logs written in the past the ``dmap`` command can be used. For example the following command aggregates all mapreduce fields of all the logs and calculates the average memory free grouped by day of the month, hour, minute and the server hostname. ``dmap`` will print interim results every few seconds. The final result however will be written to file ``mapreduce.csv``. diff --git a/doc/quickstart.md b/doc/quickstart.md index 733442f..6baedbb 100644 --- a/doc/quickstart.md +++ b/doc/quickstart.md @@ -16,7 +16,7 @@ On Linux you need to install the libacl development library for file system ACL To compile and install all DTail binaries directly from GitHub run: ```console -% for cmd in dcat dgrep dmap drun dtail dserver; do +% for cmd in dcat dgrep dmap dtail dserver; do go get github.com/mimecast/dtail/cmd/$cmd; done ``` @@ -26,7 +26,6 @@ It produces the following executables in ``$GOPATH/bin``: * ``dcat``: Client for displaying whole files remotely (distributed cat) * ``dgrep``: Client for searching whole files files remotely using a regex (distributed grep) * ``dmap``: Client for executing distributed mapreduce queries (may will consume a lot of RAM and CPU) -* ``drun``: Client for executing commands on remote servers. * ``dtail``: Client for tailing/following log files remotely (distributed tail) * ``dserver``: The DTail server diff --git a/internal/clients/runclient.go b/internal/clients/runclient.go deleted file mode 100644 index 5464d54..0000000 --- a/internal/clients/runclient.go +++ /dev/null @@ -1,87 +0,0 @@ -package clients - -import ( - "crypto/sha256" - "encoding/base64" - "encoding/hex" - "fmt" - "runtime" - "strings" - - "github.com/mimecast/dtail/internal/clients/handlers" - "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/omode" -) - -// RunClient is a client to run various commands on the server. -type RunClient struct { - baseClient - jobName string - background string -} - -// NewRunClient returns a new run client to execute commands on the remote server. -func NewRunClient(args Args, background, jobName string) (*RunClient, error) { - args.Mode = omode.RunClient - - if jobName == "" { - jobName = hash(strings.Join(args.Arguments, " ")) - } - - c := RunClient{ - baseClient: baseClient{ - Args: args, - throttleCh: make(chan struct{}, args.ConnectionsPerCPU*runtime.NumCPU()), - retry: false, - }, - jobName: jobName, - background: background, - } - - c.init() - c.makeConnections(c) - - return &c, nil -} - -func (c RunClient) makeHandler(server string) handlers.Handler { - return handlers.NewClientHandler(server) -} - -func (c RunClient) makeCommands() (commands []string) { - if c.Timeout > 0 { - commands = append(commands, fmt.Sprintf("timeout %d run%s %s", c.Timeout, c.options(), c.What)) - return - } - - commands = append(commands, fmt.Sprintf("run%s %s", c.options(), c.What)) - return -} - -func (c RunClient) options() string { - var sb strings.Builder - - logger.Debug("options", fmt.Sprintf(":background=%s", c.background)) - sb.WriteString(fmt.Sprintf(":background=%s", c.background)) - - logger.Debug("options", fmt.Sprintf(":jobName=%s", c.jobName)) - sb.WriteString(fmt.Sprintf(":jobName=%s", c.jobName)) - - if len(c.Arguments) > 0 { - logger.Debug("options", fmt.Sprintf(":outerArgs=base64%%%s", strings.Join(c.Arguments, " "))) - sb.WriteString(fmt.Sprintf(":outerArgs=base64%%%s", encode64(strings.Join(c.Arguments, " ")))) - } - - return sb.String() -} - -func encode64(str string) string { - return base64.StdEncoding.EncodeToString([]byte(str)) -} - -func hash(str string) string { - h := sha256.New() - h.Write([]byte(str)) - - return hex.EncodeToString(h.Sum(nil)) -} diff --git a/internal/io/run/run.go b/internal/io/run/run.go deleted file mode 100644 index 2bb3756..0000000 --- a/internal/io/run/run.go +++ /dev/null @@ -1,150 +0,0 @@ -package run - -import ( - "bufio" - "context" - "io" - "os/exec" - "sync" - "syscall" - "time" - - "github.com/mimecast/dtail/internal/io/line" - "github.com/mimecast/dtail/internal/io/logger" -) - -// Run is for execute a command. -type Run struct { - command string - args []string - cmd *exec.Cmd - pgroupKilled chan struct{} -} - -// New returns a new command runner. -func New(command string, args []string) Run { - return Run{ - command: command, - args: args, - pgroupKilled: make(chan struct{}), - } -} - -// StartBackground starts running the command in background. -func (r Run) StartBackground(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, lines chan<- line.Line) (pid int, err error) { - pid = -1 - - if len(r.args) > 0 { - logger.Debug(r.command, r.args, " ") - r.cmd = exec.CommandContext(ctx, r.command, r.args...) - } else { - logger.Debug(r.command) - r.cmd = exec.CommandContext(ctx, r.command) - } - - // Create a new process group, so that kill() will only kill this command + pgroup. - r.cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} - - stdoutPipe, myErr := r.cmd.StdoutPipe() - if err != nil { - wg.Done() - err = myErr - return - } - - stderrPipe, myErr := r.cmd.StderrPipe() - if myErr != nil { - wg.Done() - err = myErr - return - } - - if myErr := r.cmd.Start(); err != nil { - wg.Done() - err = myErr - return - } - - if r.cmd.Process != nil { - pid = r.cmd.Process.Pid - } - - commandExited := make(chan struct{}) - - var pipeWg sync.WaitGroup - pipeWg.Add(2) - - go r.killPgroup(ctx, commandExited, pid) - go r.pipeToLines(commandExited, &pipeWg, pid, stdoutPipe, "STDOUT", lines) - go r.pipeToLines(commandExited, &pipeWg, pid, stderrPipe, "STDERR", lines) - - go func() { - exitCode := 255 - if waitErr := r.cmd.Wait(); waitErr != nil { - if exitError, ok := waitErr.(*exec.ExitError); ok { - exitCode = exitError.ExitCode() - } - } - ec <- exitCode - - // Tell pipes we are done - close(commandExited) - // Wait for process group to be killed - <-r.pgroupKilled - // Wait for the pipes to flush the contents - pipeWg.Wait() - // Now the job is truly done - wg.Done() - }() - - return -} - -func (r Run) pipeToLines(commandExited chan struct{}, wg *sync.WaitGroup, pid int, reader io.Reader, what string, lines chan<- line.Line) { - defer wg.Done() - bufReader := bufio.NewReader(reader) - - for { - time.Sleep(time.Millisecond * 10) - lineStr, err := bufReader.ReadString('\n') - - if err != nil { - select { - case <-commandExited: - return - default: - } - continue - } - - newLine := line.Line{ - Content: []byte(lineStr), - Count: uint64(pid), - TransmittedPerc: 100, - SourceID: what, - } - - select { - case lines <- newLine: - case <-commandExited: - return - } - } -} - -func (r Run) killPgroup(ctx context.Context, commandExited chan struct{}, pid int) { - if pid == -1 { - close(r.pgroupKilled) - return - } - - if pgid, err := syscall.Getpgid(pid); err == nil { - // Kill process group when done - select { - case <-ctx.Done(): - case <-commandExited: - } - syscall.Kill(-pgid, syscall.SIGKILL) - close(r.pgroupKilled) - } -} diff --git a/internal/server/background/background.go b/internal/server/background/background.go deleted file mode 100644 index 3907448..0000000 --- a/internal/server/background/background.go +++ /dev/null @@ -1,126 +0,0 @@ -package background - -import ( - "context" - "errors" - "fmt" - "strings" - "sync" - - "github.com/mimecast/dtail/internal/io/logger" -) - -type job struct { - cancel context.CancelFunc - wg *sync.WaitGroup -} - -// Background specifies a job or command run in background on server side. -// This does not require an active DTail client SSH connection/session. -type Background struct { - mutex *sync.Mutex - jobs map[string]job -} - -// New returns a new background manager. -func New() Background { - return Background{ - jobs: make(map[string]job), - mutex: &sync.Mutex{}, - } -} - -// Add a background job. -func (b Background) Add(userName, jobName string, cancel context.CancelFunc, wg *sync.WaitGroup) error { - key := b.key(userName, jobName) - logger.Debug("background", "Add", key) - - b.mutex.Lock() - defer b.mutex.Unlock() - - if _, ok := b.jobs[key]; ok { - return errors.New("job already exists") - } - - b.jobs[key] = job{cancel, wg} - - // Clean up background job database. - go func() { - wg.Wait() - b.cancel(key) - }() - - return nil -} - -// Cancel a background job. -func (b Background) Cancel(userName, jobName string) error { - key := b.key(userName, jobName) - logger.Debug("background", "Cancel", key) - - return b.cancel(key) -} - -func (b Background) cancel(key string) error { - job, ok := b.get(key) - logger.Debug("background", "cancel", key, job, ok) - - if !ok { - return errors.New("no job to cancel") - } - - logger.Debug("background", "cancel", "run job.cancel()") - job.cancel() - logger.Debug("background", "cancel", "run job.wg.Wait()") - job.wg.Wait() - logger.Debug("background", "cancel", "run b.delete(key)") - b.delete(key) - - return nil -} - -// ListJobsC returns a channel listing all jobs of the given user. -func (b Background) ListJobsC(userName string) <-chan string { - logger.Debug("background", "ListJobC", userName) - - ch := make(chan string) - - go func() { - defer close(ch) - - b.mutex.Lock() - defer b.mutex.Unlock() - - for k := range b.jobs { - logger.Debug("ListJobsC", k, userName) - if strings.HasPrefix(k, fmt.Sprintf("%s.", userName)) { - ch <- k - } - } - }() - - return ch -} - -func (b Background) get(key string) (job, bool) { - logger.Debug("background", "get", key) - - b.mutex.Lock() - defer b.mutex.Unlock() - - job, ok := b.jobs[key] - return job, ok -} - -func (b Background) delete(key string) { - logger.Debug("background", "delete", key) - - b.mutex.Lock() - defer b.mutex.Unlock() - - delete(b.jobs, key) -} - -func (Background) key(userName, jobName string) string { - return fmt.Sprintf("%s.%s", userName, jobName) -} diff --git a/internal/server/handlers/runcommand.go b/internal/server/handlers/runcommand.go deleted file mode 100644 index 8e5895b..0000000 --- a/internal/server/handlers/runcommand.go +++ /dev/null @@ -1,111 +0,0 @@ -package handlers - -import ( - "context" - "errors" - "fmt" - "io/ioutil" - "os" - "os/exec" - "strings" - "sync" - "time" - - "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/io/run" -) - -type runCommand struct { - server *ServerHandler - run run.Run -} - -func newRunCommand(server *ServerHandler) runCommand { - return runCommand{ - server: server, - } -} - -func (r runCommand) StartBackground(ctx context.Context, wg *sync.WaitGroup, argc int, args, outerArgs []string) error { - if argc < 2 { - return fmt.Errorf("%s: args:%v argc:%d", commandParseWarning, args, argc) - } - - ec := make(chan int, 1) - var pid int - var err error - - command := strings.Join(args[1:], " ") - if strings.Contains(command, ";") || strings.Contains(command, "\n") { - if pid, err = r.startScript(ctx, wg, ec, command, outerArgs); err != nil { - r.server.sendServerMessage(".run exitstatus 255") - return err - } - return nil - } - - if pid, err = r.start(ctx, wg, ec, strings.TrimSpace(command), outerArgs); err != nil { - r.server.sendServerMessage(".run exitstatus 255") - return err - } - - exitCode := <-ec - r.server.sendServerMessage(fmt.Sprintf(".run exitstatus %d", exitCode)) - r.server.sendServerMessage(logger.Info(fmt.Sprintf("Process %d exited with status %d", pid, exitCode))) - - return nil -} - -func (r runCommand) startScript(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, script string, outerArgs []string) (int, error) { - if _, err := os.Stat(config.Common.TmpDir); os.IsNotExist(err) { - return -1, err - } - - timestamp := time.Now().UnixNano() - scriptPath := fmt.Sprintf("%s/%s_%v.sh", config.Common.TmpDir, r.server.user.Name, timestamp) - - // TODO: On dserver startup delete all previously written scripts (there might be left overs due to a crash or so) - logger.Debug(r.server.user, "Writing temp script", scriptPath) - - script = fmt.Sprintf("#!/bin/sh\n%s", script) - if err := ioutil.WriteFile(scriptPath, []byte(script), 0700); err != nil { - return -1, err - } - - pid, err := r.start(ctx, wg, ec, scriptPath, outerArgs) - go func() { - wg.Wait() - logger.Debug("Deleting script", scriptPath) - os.Remove(scriptPath) - }() - - return pid, err -} - -func (r runCommand) start(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, command string, outerArgs []string) (int, error) { - if len(command) == 0 { - return -1, errors.New("Empty command provided") - } - - splitted := strings.Split(command, " ") - path := splitted[0] - args := splitted[1:] - args = append(args, outerArgs...) - - qualifiedPath, err := exec.LookPath(path) - if err != nil { - return -1, err - } - - if !r.server.user.HasFilePermission(qualifiedPath, "runcommands") { - return -1, fmt.Errorf("No permission to execute path: %s", qualifiedPath) - } - - r.run = run.New(qualifiedPath, args) - pid, err := r.run.StartBackground(ctx, wg, ec, r.server.lines) - if err != nil { - return pid, err - } - return pid, nil -} diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 164a280..7ad1224 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -9,7 +9,6 @@ import ( "os" "strconv" "strings" - "sync" "sync/atomic" "time" @@ -19,7 +18,6 @@ import ( "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/mapr/server" "github.com/mimecast/dtail/internal/omode" - "github.com/mimecast/dtail/internal/server/background" user "github.com/mimecast/dtail/internal/user/server" "github.com/mimecast/dtail/internal/version" ) @@ -47,11 +45,10 @@ type ServerHandler struct { ackCloseReceived chan struct{} activeCommands int32 activeReaders int32 - background background.Background } // NewServerHandler returns the server handler. -func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}, background background.Background) *ServerHandler { +func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}) *ServerHandler { h := ServerHandler{ done: internal.NewDone(), lines: make(chan line.Line, 100), @@ -63,7 +60,6 @@ func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWait globalServerWaitFor: globalServerWaitFor, regex: ".", user: user, - background: background, } fqdn, err := os.Hostname() @@ -314,85 +310,6 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args [] commandFinished() }() - case "run": - // TODO: Refactor this "run" case, move code to runcommand.go - command := newRunCommand(h) - - jobName, _ := options["jobName"] - logger.Debug(h.user, "run", options) - - if val, ok := options["background"]; ok && (val == "cancel" || val == "stop") { - if err := h.background.Cancel(h.user.Name, jobName); err != nil { - h.sendServerMessage(logger.Error(h.user, err, jobName, args)) - } else { - h.sendServerMessage(logger.Info(h.user, "job cancelled", jobName)) - } - commandFinished() - return - } - - if val, ok := options["background"]; ok && val == "list" { - h.sendServerMessage("Listing jobs") - count := 0 - for jobName := range h.background.ListJobsC(h.user.Name) { - h.sendServerMessage(jobName) - count++ - } - h.sendServerMessage(fmt.Sprintf("Found %d jobs", count)) - commandFinished() - return - } - - str, _ := options["outerArgs"] - outerArgs := strings.Split(str, " ") - - var background bool - if val, ok := options["background"]; ok && val == "start" { - background = true - } - - var wg sync.WaitGroup - wg.Add(1) - - if background { - if timeout == 0 { - // Set default background timeout. - timeout = time.Hour * 1 - } - - commandCtx, cancel := context.WithTimeout(ctx, timeout) - - if err := h.background.Add(h.user.Name, jobName, cancel, &wg); err != nil { - h.sendServerMessage(logger.Error(h.user, err, jobName, args)) - commandFinished() - return - } - ctx = commandCtx - } - - if err := command.StartBackground(ctx, &wg, argc, args, outerArgs); err != nil { - h.sendServerMessage(logger.Error(h.user, "Unable to execute command", argc, args, err)) - commandFinished() - return - } - - // Make sure that server waits for all sub-processes to finish on shutdown - go func() { h.globalServerWaitFor <- struct{}{} }() - go func() { - wg.Wait() - <-h.globalServerWaitFor - }() - - if background { - h.sendServerMessage(logger.Info(h.user, jobName, "job started in background")) - commandFinished() - return - } - - // Command run in foreground, wait for it to complete before finishing the connection. - wg.Wait() - commandFinished() - case "ack", ".ack": h.handleAckCommand(argc, args) commandFinished() diff --git a/internal/server/server.go b/internal/server/server.go index 5e2a521..d4255a3 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -11,7 +11,6 @@ import ( "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/logger" - "github.com/mimecast/dtail/internal/server/background" "github.com/mimecast/dtail/internal/server/handlers" "github.com/mimecast/dtail/internal/ssh/server" user "github.com/mimecast/dtail/internal/user/server" @@ -36,8 +35,6 @@ type Server struct { cont *continuous // Wait counter, e.g. there might be still subprocesses (forked by drun) to be killed. shutdownWaitFor chan struct{} - // Background jobs - background background.Background } // New returns a new server. @@ -51,7 +48,6 @@ func New() *Server { shutdownWaitFor: make(chan struct{}, 1000), sched: newScheduler(), cont: newContinuous(), - background: background.New(), } s.sshServerConfig.PasswordCallback = s.Callback @@ -183,7 +179,7 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch case config.ControlUser: handler = handlers.NewControlHandler(user) default: - handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor, s.background) + handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor) } terminate := func() { -- cgit v1.2.3 From 1626b6237d87a6ae6238b0d297d7bfb5cf7ddd38 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 16 Nov 2020 22:16:10 +0000 Subject: strip out timeout which was used for background command scheduling --- internal/server/handlers/serverhandler.go | 34 ++----------------------------- 1 file changed, 2 insertions(+), 32 deletions(-) diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go index 7ad1224..843eabc 100644 --- a/internal/server/handlers/serverhandler.go +++ b/internal/server/handlers/serverhandler.go @@ -7,7 +7,6 @@ import ( "fmt" "io" "os" - "strconv" "strings" "sync/atomic" "time" @@ -147,7 +146,6 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) { func (h *ServerHandler) handleCommand(commandStr string) { logger.Debug(h.user, commandStr) - var timeout time.Duration ctx := context.Background() args, argc, err := h.handleProtocolVersion(strings.Split(commandStr, " ")) @@ -162,12 +160,6 @@ func (h *ServerHandler) handleCommand(commandStr string) { return } - args, argc, timeout, err = h.handleTimeout(args, argc) - if err != nil { - h.send(h.serverMessages, logger.Error(h.user, err)) - return - } - if h.user.Name == config.ControlUser { h.handleControlCommand(argc, args) return @@ -179,19 +171,7 @@ func (h *ServerHandler) handleCommand(commandStr string) { cancel() }() - if timeout > 0 { - logger.Info(h.user, "Command with timeout context", argc, args, timeout) - ctx, cancel := context.WithTimeout(ctx, timeout) - go func() { - <-ctx.Done() - logger.Info(h.user, "Command timed out, canceling it", args, args, timeout) - cancel() - }() - h.handleUserCommand(ctx, argc, args, timeout) - return - } - - h.handleUserCommand(ctx, argc, args, timeout) + h.handleUserCommand(ctx, argc, args) } func (h *ServerHandler) handleProtocolVersion(args []string) ([]string, int, error) { @@ -229,16 +209,6 @@ func (h *ServerHandler) handleBase64(args []string, argc int) ([]string, int, er return args, argc, nil } -func (h *ServerHandler) handleTimeout(args []string, argc int) ([]string, int, tim