From ac9c011475e4aefeb4ef70649e0977e6a3b370fb Mon Sep 17 00:00:00 2001 From: Tom Wilkie Date: Tue, 10 Nov 2015 13:46:31 +0000 Subject: [PATCH] Pipe plumbing - Add store of pipes in the app - Add pipe type, handling impedance mismatch, used in app and probe. - App <-> Probe pipes have their own websockets. - Add pipe websocket endpoint in app. - Pipe IDs are strings, lose the request/response IDs, and give the json encoder lowercase field names. - Add simple golang ws client, for testing. - Pipe lifecycle plumbing. - Ref count and timeout both ends of pipes in the app - Deal with POST /api/pipe/:pid?_method=delete - Add end-to-end unit test for pipes. - Add test for timing out pipes. - Update go-docker client to tomwilkie/go-dockerclient - Backend work for non-raw ttys - Close pipes when they close themselves in the probe - Ensure all http connections are done before returning from client.Stop() --- app/controls.go | 16 +-- app/controls_test.go | 22 ++-- app/pipes.go | 194 +++++++++++++++++++++++++++++ app/pipes_internal_test.go | 140 +++++++++++++++++++++ experimental/demoprobe/main.go | 2 +- experimental/fixprobe/main.go | 2 +- experimental/pipeclient/.gitignore | 1 + experimental/pipeclient/main.go | 77 ++++++++++++ probe/controls/controls.go | 10 +- probe/controls/controls_test.go | 4 - probe/controls/pipes.go | 48 +++++++ prog/app.go | 1 + prog/probe.go | 11 +- xfer/app_client.go | 160 ++++++++++++++++++------ xfer/app_client_internal_test.go | 6 +- xfer/controls.go | 73 ++++++----- xfer/multi_client.go | 39 ++++-- xfer/multi_client_test.go | 14 +-- xfer/pipes.go | 128 +++++++++++++++++++ 19 files changed, 829 insertions(+), 119 deletions(-) create mode 100644 app/pipes.go create mode 100644 app/pipes_internal_test.go create mode 100644 experimental/pipeclient/.gitignore create mode 100644 experimental/pipeclient/main.go create mode 100644 probe/controls/pipes.go create mode 100644 xfer/pipes.go diff --git a/app/controls.go b/app/controls.go index f2f5f94fc..c7bc4518d 100644 --- a/app/controls.go +++ b/app/controls.go @@ -24,6 +24,12 @@ func RegisterControlRoutes(router *mux.Router) { type controlHandler struct { id int64 client *rpc.Client + codec *xfer.JSONWebsocketCodec +} + +type controlRouter struct { + sync.Mutex + probes map[string]controlHandler } func (ch *controlHandler) handle(req xfer.Request) xfer.Response { @@ -34,11 +40,6 @@ func (ch *controlHandler) handle(req xfer.Request) xfer.Response { return res } -type controlRouter struct { - sync.Mutex - probes map[string]controlHandler -} - func (cr *controlRouter) get(probeID string) (controlHandler, bool) { cr.Lock() defer cr.Unlock() @@ -79,7 +80,7 @@ func (cr *controlRouter) handleControl(w http.ResponseWriter, r *http.Request) { } result := handler.handle(xfer.Request{ - ID: rand.Int63(), + AppID: UniqueID, NodeID: nodeID, Control: control, }) @@ -87,7 +88,7 @@ func (cr *controlRouter) handleControl(w http.ResponseWriter, r *http.Request) { respondWith(w, http.StatusBadRequest, result.Error) return } - respondWith(w, http.StatusOK, result.Value) + respondWith(w, http.StatusOK, result) } // handleProbeWS accepts websocket connections from the probe and registers @@ -110,6 +111,7 @@ func (cr *controlRouter) handleProbeWS(w http.ResponseWriter, r *http.Request) { client := rpc.NewClientWithCodec(codec) handler := controlHandler{ id: rand.Int63(), + codec: codec, client: client, } diff --git a/app/controls_test.go b/app/controls_test.go index 258677092..82e51bf64 100644 --- a/app/controls_test.go +++ b/app/controls_test.go @@ -29,13 +29,7 @@ func TestControl(t *testing.T) { probeConfig := xfer.ProbeConfig{ ProbeID: "foo", } - client, err := xfer.NewAppClient(probeConfig, ip+":"+port, ip+":"+port) - if err != nil { - t.Fatal(err) - } - defer client.Stop() - - client.ControlConnection(xfer.ControlHandlerFunc(func(req xfer.Request) xfer.Response { + controlHandler := xfer.ControlHandlerFunc(func(req xfer.Request) xfer.Response { if req.NodeID != "nodeid" { t.Fatalf("'%s' != 'nodeid'", req.NodeID) } @@ -47,7 +41,13 @@ func TestControl(t *testing.T) { return xfer.Response{ Value: "foo", } - })) + }) + client, err := xfer.NewAppClient(probeConfig, ip+":"+port, ip+":"+port, controlHandler) + if err != nil { + t.Fatal(err) + } + client.ControlConnection() + defer client.Stop() time.Sleep(100 * time.Millisecond) @@ -59,12 +59,12 @@ func TestControl(t *testing.T) { t.Fatal(err) } - var response string + var response xfer.Response if err := json.NewDecoder(resp.Body).Decode(&response); err != nil { t.Fatal(err) } - if response != "foo" { - t.Fatalf("'%s' != 'foo'", response) + if response.Value != "foo" { + t.Fatalf("'%s' != 'foo'", response.Value) } } diff --git a/app/pipes.go b/app/pipes.go new file mode 100644 index 000000000..75797a56b --- /dev/null +++ b/app/pipes.go @@ -0,0 +1,194 @@ +package app + +import ( + "io" + "log" + "net/http" + "sync" + "time" + + "github.com/gorilla/mux" + + "github.com/weaveworks/scope/common/mtime" + "github.com/weaveworks/scope/xfer" +) + +const ( + gcInterval = 30 * time.Second // we check all the pipes every 30s + pipeTimeout = 1 * time.Minute // pipes are closed when a client hasn't been connected for 1 minute + gcTimeout = 1 * time.Minute // after another 1 minute, tombstoned pipes are forgotten +) + +// PipeRouter connects incoming and outgoing pipes. +type PipeRouter struct { + sync.Mutex + wait sync.WaitGroup + quit chan struct{} + pipes map[string]*pipe +} + +// for each end of the pipe, we keep a reference count & lastUsedTIme, +// such that we can timeout pipes when either end is inactive. +type end struct { + refCount int + lastUsedTime time.Time +} + +type pipe struct { + ui, probe end + tombstoneTime time.Time + + xfer.Pipe +} + +// RegisterPipeRoutes registers the pipe routes +func RegisterPipeRoutes(router *mux.Router) *PipeRouter { + pipeRouter := &PipeRouter{ + quit: make(chan struct{}), + pipes: map[string]*pipe{}, + } + pipeRouter.wait.Add(1) + go pipeRouter.gcLoop() + router.Methods("GET").Path("/api/pipe/{pipeID}").HandlerFunc(pipeRouter.handleWs(func(p *pipe) (*end, io.ReadWriter) { + uiEnd, _ := p.Ends() + return &p.ui, uiEnd + })) + router.Methods("GET").Path("/api/pipe/{pipeID}/probe").HandlerFunc(pipeRouter.handleWs(func(p *pipe) (*end, io.ReadWriter) { + _, probeEnd := p.Ends() + return &p.probe, probeEnd + })) + router.Methods("DELETE").Path("/api/pipe/{pipeID}").HandlerFunc(pipeRouter.deletePipe) + router.Methods("POST").Path("/api/pipe/{pipeID}").HandlerFunc(pipeRouter.deletePipe) + return pipeRouter +} + +// Stop stops the pipeRouter +func (pr *PipeRouter) Stop() { + close(pr.quit) + pr.wait.Wait() +} + +func (pr *PipeRouter) gcLoop() { + defer pr.wait.Done() + ticker := time.Tick(gcInterval) + for { + select { + case <-pr.quit: + return + case <-ticker: + } + + pr.timeoutPipes() + pr.gcPipes() + } +} + +func (pr *PipeRouter) timeoutPipes() { + pr.Lock() + defer pr.Unlock() + now := mtime.Now() + for id, pipe := range pr.pipes { + if pipe.Closed() || (pipe.ui.refCount > 0 && pipe.probe.refCount > 0) { + continue + } + + if (pipe.ui.refCount == 0 && now.Sub(pipe.ui.lastUsedTime) >= pipeTimeout) || + (pipe.probe.refCount == 0 && now.Sub(pipe.probe.lastUsedTime) >= pipeTimeout) { + log.Printf("Timing out pipe %s", id) + pipe.Close() + pipe.tombstoneTime = now + } + } +} + +func (pr *PipeRouter) gcPipes() { + pr.Lock() + defer pr.Unlock() + now := mtime.Now() + for pipeID, pipe := range pr.pipes { + if pipe.Closed() && now.Sub(pipe.tombstoneTime) >= gcTimeout { + delete(pr.pipes, pipeID) + } + } +} + +func (pr *PipeRouter) getOrCreatePipe(id string) (*pipe, bool) { + pr.Lock() + defer pr.Unlock() + p, ok := pr.pipes[id] + if !ok { + log.Printf("Creating pipe id %s", id) + p = &pipe{ + ui: end{lastUsedTime: mtime.Now()}, + probe: end{lastUsedTime: mtime.Now()}, + Pipe: xfer.NewPipe(), + } + pr.pipes[id] = p + } + if p.Closed() { + return nil, false + } + return p, true +} + +func (pr *PipeRouter) getPipeRef(id string, pipe *pipe, end *end) bool { + pr.Lock() + defer pr.Unlock() + if pipe.Closed() { + return false + } + end.refCount++ + return true +} + +func (pr *PipeRouter) putPipeRef(id string, pipe *pipe, end *end) { + pr.Lock() + defer pr.Unlock() + + end.refCount-- + if end.refCount != 0 { + return + } + + if !pipe.Closed() { + end.lastUsedTime = mtime.Now() + } +} + +func (pr *PipeRouter) handleWs(endSelector func(*pipe) (*end, io.ReadWriter)) func(http.ResponseWriter, *http.Request) { + return func(w http.ResponseWriter, r *http.Request) { + pipeID := mux.Vars(r)["pipeID"] + pipe, ok := pr.getOrCreatePipe(pipeID) + if !ok { + http.NotFound(w, r) + return + } + + endRef, endIO := endSelector(pipe) + if !pr.getPipeRef(pipeID, pipe, endRef) { + http.NotFound(w, r) + return + } + defer pr.putPipeRef(pipeID, pipe, endRef) + + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + log.Printf("Error upgrading to websocket: %v", err) + return + } + defer conn.Close() + + pipe.CopyToWebsocket(endIO, conn) + } +} + +func (pr *PipeRouter) deletePipe(w http.ResponseWriter, r *http.Request) { + pipeID := mux.Vars(r)["pipeID"] + pipe, ok := pr.getOrCreatePipe(pipeID) + if ok && pr.getPipeRef(pipeID, pipe, &pipe.ui) { + log.Printf("Closing pipe %s", pipeID) + pipe.Close() + pipe.tombstoneTime = mtime.Now() + pr.putPipeRef(pipeID, pipe, &pipe.ui) + } +} diff --git a/app/pipes_internal_test.go b/app/pipes_internal_test.go new file mode 100644 index 000000000..aad796ff0 --- /dev/null +++ b/app/pipes_internal_test.go @@ -0,0 +1,140 @@ +package app + +import ( + "bytes" + "fmt" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/mux" + "github.com/gorilla/websocket" + + "github.com/weaveworks/scope/common/mtime" + "github.com/weaveworks/scope/probe/controls" + "github.com/weaveworks/scope/test" + "github.com/weaveworks/scope/xfer" +) + +func TestPipeTimeout(t *testing.T) { + router := mux.NewRouter() + pr := RegisterPipeRoutes(router) + pr.Stop() // we don't want the loop running in the background + + mtime.NowForce(time.Now()) + defer mtime.NowReset() + + // create a new pipe. + id := "foo" + pipe, ok := pr.getOrCreatePipe(id) + if !ok { + t.Fatalf("not ok") + } + + // move time forward such that the new pipe should timeout + mtime.NowForce(mtime.Now().Add(pipeTimeout)) + pr.timeoutPipes() + if !pipe.Closed() { + t.Fatalf("pipe didn't timeout") + } + + // move time forward such that the pipe should be GCd + mtime.NowForce(mtime.Now().Add(gcTimeout)) + pr.gcPipes() + if _, ok := pr.pipes[id]; ok { + t.Fatalf("pipe not gc'd") + } +} + +type adapter struct { + c xfer.AppClient +} + +func (a adapter) PipeConnection(_, pipeID string, pipe xfer.Pipe) error { + a.c.PipeConnection(pipeID, pipe) + return nil +} + +func (a adapter) PipeClose(_, pipeID string) error { + return a.c.PipeClose(pipeID) +} + +func TestPipeClose(t *testing.T) { + router := mux.NewRouter() + pr := RegisterPipeRoutes(router) + defer pr.Stop() + + server := httptest.NewServer(router) + defer server.Close() + + ip, port, err := net.SplitHostPort(strings.TrimPrefix(server.URL, "http://")) + if err != nil { + t.Fatal(err) + } + + probeConfig := xfer.ProbeConfig{ + ProbeID: "foo", + } + client, err := xfer.NewAppClient(probeConfig, ip+":"+port, ip+":"+port, nil) + if err != nil { + t.Fatal(err) + } + defer client.Stop() + + oldClient := controls.Client + defer func() { controls.Client = oldClient }() + controls.Client = adapter{client} + + // this is the probe end of the pipe + pipeID, pipe, err := controls.NewPipe("appid") + if err != nil { + t.Fatal(err) + } + + // this is a client to the app + pipeURL := fmt.Sprintf("ws://%s:%s/api/pipe/%s", ip, port, pipeID) + conn, _, err := websocket.DefaultDialer.Dial(pipeURL, http.Header{}) + if err != nil { + t.Fatal(err) + } + + // Send something from pipe -> app -> conn + local, _ := pipe.Ends() + msg := []byte("hello world") + if _, err := local.Write(msg); err != nil { + t.Fatal(err) + } + + if _, buf, err := conn.ReadMessage(); err != nil { + t.Fatal(err) + } else if !bytes.Equal(buf, msg) { + t.Fatalf("%v != %v", buf, msg) + } + + // Send something from conn -> app -> probe + msg = []byte("goodbye, cruel world") + if err := conn.WriteMessage(websocket.BinaryMessage, msg); err != nil { + t.Fatal(err) + } + + buf := make([]byte, 1024) + if n, err := local.Read(buf); err != nil { + t.Fatal(err) + } else if !bytes.Equal(msg, buf[:n]) { + t.Fatalf("%v != %v", buf, msg) + } + + // Now delete the pipe + if err := pipe.Close(); err != nil { + t.Fatal(err) + } + + // the client backs off for 1 second before trying to reconnect the pipe, + // so we need to wait for longer. + test.Poll(t, 2*time.Second, true, func() interface{} { + return pipe.Closed() + }) +} diff --git a/experimental/demoprobe/main.go b/experimental/demoprobe/main.go index 613966f2a..57dbc8e4c 100644 --- a/experimental/demoprobe/main.go +++ b/experimental/demoprobe/main.go @@ -27,7 +27,7 @@ func main() { Token: "demoprobe", ProbeID: "demoprobe", Insecure: false, - }, *publish, *publish) + }, *publish, *publish, nil) if err != nil { log.Fatal(err) } diff --git a/experimental/fixprobe/main.go b/experimental/fixprobe/main.go index 78a70d442..d19eb99b4 100644 --- a/experimental/fixprobe/main.go +++ b/experimental/fixprobe/main.go @@ -38,7 +38,7 @@ func main() { Token: "fixprobe", ProbeID: "fixprobe", Insecure: false, - }, *publish, *publish) + }, *publish, *publish, nil) if err != nil { log.Fatal(err) } diff --git a/experimental/pipeclient/.gitignore b/experimental/pipeclient/.gitignore new file mode 100644 index 000000000..3ba9ee9a6 --- /dev/null +++ b/experimental/pipeclient/.gitignore @@ -0,0 +1 @@ +pipeclient diff --git a/experimental/pipeclient/main.go b/experimental/pipeclient/main.go new file mode 100644 index 000000000..bf986bae9 --- /dev/null +++ b/experimental/pipeclient/main.go @@ -0,0 +1,77 @@ +package main + +import ( + "log" + "os" + + "github.com/gorilla/websocket" + //"golang.org/x/crypto/ssh/terminal" + "github.com/davecgh/go-spew/spew" +) + +func main() { + if len(os.Args) < 2 { + log.Fatal("must specify url") + } + url := os.Args[1] + log.Printf("Connecting to %s", url) + + //oldState, err := terminal.MakeRaw(0) + //if err != nil { + // panic(err) + //} + //defer terminal.Restore(0, oldState) + + dialer := websocket.Dialer{} + conn, _, err := dialer.Dial(url, nil) + if err != nil { + panic(err) + } + defer conn.Close() + readQuit := make(chan struct{}) + writeQuit := make(chan struct{}) + + // Read-from-UI loop + go func() { + defer close(readQuit) + for { + _, buf, err := conn.ReadMessage() // TODO type should be binary message + if err != nil { + log.Printf("Error reading websocket: %v", err) + return + } + + spew.Dump(buf) + + if _, err := os.Stdout.Write(buf); err != nil { + log.Printf("Error writing stdout: %v", err) + return + } + } + }() + + // Write-to-UI loop + go func() { + defer close(writeQuit) + buf := make([]byte, 1024) + for { + n, err := os.Stdin.Read(buf) + if err != nil { + log.Printf("Error reading stdin: %v", err) + return + } + + if err := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil { + log.Printf("Error writing websocket: %v", err) + return + } + } + }() + + // block until one of the goroutines exits + // this convoluted mechanism is to ensure we only close the websocket once. + select { + case <-readQuit: + case <-writeQuit: + } +} diff --git a/probe/controls/controls.go b/probe/controls/controls.go index 3dcce5d49..eae10eb69 100644 --- a/probe/controls/controls.go +++ b/probe/controls/controls.go @@ -1,7 +1,6 @@ package controls import ( - "fmt" "sync" "github.com/weaveworks/scope/xfer" @@ -18,15 +17,10 @@ func HandleControlRequest(req xfer.Request) xfer.Response { handler, ok := handlers[req.Control] mtx.Unlock() if !ok { - return xfer.Response{ - ID: req.ID, - Error: fmt.Sprintf("Control '%s' not recognised", req.Control), - } + return xfer.ResponseErrorf("Control '%s' not recognised", req.Control) } - response := handler(req) - response.ID = req.ID - return response + return handler(req) } // Register a new control handler under a given id. diff --git a/probe/controls/controls_test.go b/probe/controls/controls_test.go index b47a2f3df..17dd93f7b 100644 --- a/probe/controls/controls_test.go +++ b/probe/controls/controls_test.go @@ -18,11 +18,9 @@ func TestControls(t *testing.T) { defer controls.Rm("foo") want := xfer.Response{ - ID: 1234, Value: "bar", } have := controls.HandleControlRequest(xfer.Request{ - ID: 1234, Control: "foo", }) if !reflect.DeepEqual(want, have) { @@ -32,11 +30,9 @@ func TestControls(t *testing.T) { func TestControlsNotFound(t *testing.T) { want := xfer.Response{ - ID: 3456, Error: "Control 'baz' not recognised", } have := controls.HandleControlRequest(xfer.Request{ - ID: 3456, Control: "baz", }) if !reflect.DeepEqual(want, have) { diff --git a/probe/controls/pipes.go b/probe/controls/pipes.go new file mode 100644 index 000000000..ef281f49c --- /dev/null +++ b/probe/controls/pipes.go @@ -0,0 +1,48 @@ +package controls + +import ( + "fmt" + "math/rand" + + "github.com/weaveworks/scope/xfer" +) + +// Client is the thing the probe uses to make pipe connections. +var Client interface { + PipeConnection(string, string, xfer.Pipe) error + PipeClose(string, string) error +} + +// Pipe the probe-local type for a pipe, extending +// xfer.Pipe with the appID and a custom closer method. +type Pipe interface { + xfer.Pipe +} + +type pipe struct { + xfer.Pipe + id, appID string +} + +// NewPipe creats a new pipe and connects it to the app. +var NewPipe = func(appID string) (string, Pipe, error) { + pipeID := fmt.Sprintf("pipe-%d", rand.Int63()) + pipe := &pipe{ + Pipe: xfer.NewPipe(), + appID: appID, + id: pipeID, + } + if err := Client.PipeConnection(appID, pipeID, pipe.Pipe); err != nil { + return "", nil, err + } + return pipeID, pipe, nil +} + +func (p *pipe) Close() error { + err1 := p.Pipe.Close() + err2 := Client.PipeClose(p.appID, p.id) + if err1 != nil { + return err1 + } + return err2 +} diff --git a/prog/app.go b/prog/app.go index 415932057..5675a7865 100644 --- a/prog/app.go +++ b/prog/app.go @@ -24,6 +24,7 @@ func router(c app.Collector) *mux.Router { app.RegisterTopologyRoutes(c, router) app.RegisterReportPostHandler(c, router) app.RegisterControlRoutes(router) + app.RegisterPipeRoutes(router) router.Methods("GET").PathPrefix("/").Handler(http.FileServer(FS(false))) return router } diff --git a/prog/probe.go b/prog/probe.go index d289e4020..da19ea684 100644 --- a/prog/probe.go +++ b/prog/probe.go @@ -102,12 +102,19 @@ func probeMain() { } log.Printf("publishing to: %s", strings.Join(targets, ", ")) - clients := xfer.NewMultiAppClient(xfer.ProbeConfig{ + probeConfig := xfer.ProbeConfig{ Token: *token, ProbeID: probeID, Insecure: *insecure, - }, xfer.ControlHandlerFunc(controls.HandleControlRequest), xfer.NewAppClient) + } + clients := xfer.NewMultiAppClient(func(hostname, endpoint string) (xfer.AppClient, error) { + return xfer.NewAppClient( + probeConfig, hostname, endpoint, + xfer.ControlHandlerFunc(controls.HandleControlRequest), + ) + }) defer clients.Stop() + controls.Client = clients resolver := xfer.NewStaticResolver(targets, clients.Set) defer resolver.Stop() diff --git a/xfer/app_client.go b/xfer/app_client.go index a2305945a..f2bf570d4 100644 --- a/xfer/app_client.go +++ b/xfer/app_client.go @@ -30,58 +30,98 @@ type Details struct { // AppClient is a client to an app for dealing with controls. type AppClient interface { Details() (Details, error) - ControlConnection(handler ControlHandler) + ControlConnection() + PipeConnection(string, Pipe) + PipeClose(string) error Publish(r io.Reader) error Stop() } +// appClient is a client to an app for dealing with controls. type appClient struct { ProbeConfig - quit chan struct{} - target string - insecure bool - client http.Client + quit chan struct{} + mtx sync.Mutex + target string + client http.Client + + // Track ongoing websocket connections + conns map[string]*websocket.Conn // For publish publishLoop sync.Once + publishWait sync.WaitGroup readers chan io.Reader // For controls - controlServerCodecMtx sync.Mutex - controlServerCodec rpc.ServerCodec + control ControlHandler } -// NewAppClient makes a new AppClient. -func NewAppClient(pc ProbeConfig, hostname, target string) (AppClient, error) { +// NewAppClient makes a new appClient. +func NewAppClient(pc ProbeConfig, hostname, target string, control ControlHandler) (AppClient, error) { httpTransport, err := pc.getHTTPTransport(hostname) if err != nil { return nil, err } - appClient := &appClient{ + return &appClient{ ProbeConfig: pc, quit: make(chan struct{}), - readers: make(chan io.Reader), target: target, client: http.Client{ Transport: httpTransport, }, - } + conns: map[string]*websocket.Conn{}, + readers: make(chan io.Reader), + control: control, + }, nil +} - return appClient, nil +func (c *appClient) hasQuit() bool { + select { + case <-c.quit: + return true + default: + return false + } +} + +func (c *appClient) registerConn(id string, conn *websocket.Conn) bool { + c.mtx.Lock() + defer c.mtx.Unlock() + if c.hasQuit() { + conn.Close() + return false + } + c.conns[id] = conn + return true +} + +func (c *appClient) closeConn(id string) { + c.mtx.Lock() + defer c.mtx.Unlock() + if conn, ok := c.conns[id]; ok { + conn.Close() + delete(c.conns, id) + } } // Stop stops the appClient. func (c *appClient) Stop() { - c.controlServerCodecMtx.Lock() - defer c.controlServerCodecMtx.Unlock() + c.mtx.Lock() + defer c.mtx.Unlock() close(c.readers) close(c.quit) - if c.controlServerCodec != nil { - c.controlServerCodec.Close() + + for _, conn := range c.conns { + conn.Close() } + c.conns = map[string]*websocket.Conn{} c.client.Transport.(*http.Transport).CloseIdleConnections() + + c.publishWait.Wait() + return } // Details fetches the details (version, id) of the app. @@ -125,7 +165,7 @@ func (c *appClient) doWithBackoff(msg string, f func() (bool, error)) { } } -func (c *appClient) controlConnection(handler ControlHandler) error { +func (c *appClient) controlConnection() error { dialer := websocket.Dialer{} headers := http.Header{} c.ProbeConfig.authorizeHeaders(headers) @@ -142,36 +182,26 @@ func (c *appClient) controlConnection(handler ControlHandler) error { codec := NewJSONWebsocketCodec(conn) server := rpc.NewServer() - if err := server.RegisterName("control", handler); err != nil { + if err := server.RegisterName("control", c.control); err != nil { return err } - c.controlServerCodecMtx.Lock() - c.controlServerCodec = codec - // At this point we may have tried to quit earlier, so check to see if the - // quit channel has been closed, non-blocking. - select { - default: - case <-c.quit: - codec.Close() + // Will return false if we are exiting + if !c.registerConn("control", conn) { return nil } - c.controlServerCodecMtx.Unlock() + defer c.closeConn("control") server.ServeCodec(codec) - - c.controlServerCodecMtx.Lock() - c.controlServerCodec = nil - c.controlServerCodecMtx.Unlock() return nil } -func (c *appClient) ControlConnection(handler ControlHandler) { +func (c *appClient) ControlConnection() { go func() { log.Printf("Control connection to %s starting", c.target) defer log.Printf("Control connection to %s exiting", c.target) c.doWithBackoff("controls", func() (bool, error) { - return false, c.controlConnection(handler) + return false, c.controlConnection() }) }() } @@ -184,7 +214,6 @@ func (c *appClient) publish(r io.Reader) error { } req.Header.Set("Content-Encoding", "gzip") // req.Header.Set("Content-Type", "application/binary") // TODO: we should use http.DetectContentType(..) on the gob'ed - resp, err := c.client.Do(req) if err != nil { return err @@ -201,6 +230,16 @@ func (c *appClient) startPublishing() { go func() { log.Printf("Publish loop for %s starting", c.target) defer log.Printf("Publish loop for %s exiting", c.target) + + c.mtx.Lock() + if c.hasQuit() { + c.mtx.Unlock() + return + } + c.publishWait.Add(1) + defer c.publishWait.Done() + c.mtx.Unlock() + c.doWithBackoff("publish", func() (bool, error) { r := <-c.readers if r == nil { @@ -221,3 +260,54 @@ func (c *appClient) Publish(r io.Reader) error { } return nil } + +func (c *appClient) pipeConnection(id string, pipe Pipe) (bool, error) { + dialer := websocket.Dialer{} + headers := http.Header{} + c.ProbeConfig.authorizeHeaders(headers) + // TODO(twilkie) need to update sanitize to work with wss + url := sanitize.URL("ws://", 0, fmt.Sprintf("/api/pipe/%s/probe", id))(c.target) + conn, resp, err := dialer.Dial(url, headers) + if resp != nil && resp.StatusCode == http.StatusNotFound { + // Special handling - 404 means the app/user has closed the pipe + pipe.Close() + return true, nil + } + if err != nil { + return false, err + } + + // Will return false if we are exiting + if !c.registerConn(id, conn) { + return true, nil + } + defer c.closeConn(id) + + _, remote := pipe.Ends() + return false, pipe.CopyToWebsocket(remote, conn) +} + +func (c *appClient) PipeConnection(id string, pipe Pipe) { + go func() { + log.Printf("Pipe %s connection to %s starting", id, c.target) + defer log.Printf("Pipe %s connection to %s exiting", id, c.target) + c.doWithBackoff(id, func() (bool, error) { + return c.pipeConnection(id, pipe) + }) + }() +} + +// PipeClose closes the given pipe id on the app. +func (c *appClient) PipeClose(id string) error { + url := sanitize.URL("", 0, fmt.Sprintf("/api/pipe/%s", id))(c.target) + req, err := c.ProbeConfig.authorizedRequest("DELETE", url, nil) + if err != nil { + return err + } + resp, err := c.client.Do(req) + if err != nil { + return err + } + resp.Body.Close() + return nil +} diff --git a/xfer/app_client_internal_test.go b/xfer/app_client_internal_test.go index 4b12eb16a..d13b1d7cc 100644 --- a/xfer/app_client_internal_test.go +++ b/xfer/app_client_internal_test.go @@ -87,7 +87,7 @@ func TestAppClientPublishInternal(t *testing.T) { Insecure: false, } - p, err := NewAppClient(pc, u.Host, s.URL) + p, err := NewAppClient(pc, u.Host, s.URL, nil) if err != nil { t.Fatal(err) } @@ -126,7 +126,7 @@ func TestAppClientDetails(t *testing.T) { Insecure: false, } - p, err := NewAppClient(pc, u.Host, s.URL) + p, err := NewAppClient(pc, u.Host, s.URL, nil) if err != nil { t.Fatal(err) } @@ -173,7 +173,7 @@ func TestAppClientPublish(t *testing.T) { ProbeID: "", Insecure: false, } - p, err := NewAppClient(pc, u.Host, s.URL) + p, err := NewAppClient(pc, u.Host, s.URL, nil) if err != nil { t.Fatal(err) } diff --git a/xfer/controls.go b/xfer/controls.go index d13a4453f..59a1ec774 100644 --- a/xfer/controls.go +++ b/xfer/controls.go @@ -10,16 +10,24 @@ import ( // Request is the UI -> App -> Probe message type for control RPCs type Request struct { - ID int64 + AppID string NodeID string Control string } // Response is the Probe -> App -> UI message type for the control RPCs. type Response struct { - ID int64 - Value interface{} - Error string + Value interface{} `json:"value,omitempty"` + Error string `json:"error,omitempty"` + Pipe string `json:"pipe,omitempty"` + RawTTY bool `json:"raw_tty,omitempty"` +} + +// Message is the unions of Request, Response and PipeIO +type Message struct { + Request *rpc.Request + Response *rpc.Response + Value interface{} } // ControlHandler is interface used in the app and the probe to represent @@ -82,27 +90,44 @@ func (j *JSONWebsocketCodec) WriteRequest(r *rpc.Request, v interface{}) error { j.Lock() defer j.Unlock() - if err := j.conn.WriteJSON(r); err != nil { + if err := j.conn.WriteJSON(Message{Request: r}); err != nil { return err } - return j.conn.WriteJSON(v) + return j.conn.WriteJSON(Message{Value: v}) +} + +// WriteResponse implements rpc.ServerCodec +func (j *JSONWebsocketCodec) WriteResponse(r *rpc.Response, v interface{}) error { + j.Lock() + defer j.Unlock() + + if err := j.conn.WriteJSON(Message{Response: r}); err != nil { + return err + } + return j.conn.WriteJSON(Message{Value: v}) +} + +func (j *JSONWebsocketCodec) readMessage(v interface{}) (*Message, error) { + m := Message{Value: v} + if err := j.conn.ReadJSON(&m); err != nil { + close(j.err) + return nil, err + } + return &m, nil } // ReadResponseHeader implements rpc.ClientCodec func (j *JSONWebsocketCodec) ReadResponseHeader(r *rpc.Response) error { - err := j.conn.ReadJSON(r) - if err != nil { - close(j.err) + m, err := j.readMessage(nil) + if err == nil { + *r = *m.Response } return err } // ReadResponseBody implements rpc.ClientCodec func (j *JSONWebsocketCodec) ReadResponseBody(v interface{}) error { - err := j.conn.ReadJSON(v) - if err != nil { - close(j.err) - } + _, err := j.readMessage(v) return err } @@ -113,29 +138,15 @@ func (j *JSONWebsocketCodec) Close() error { // ReadRequestHeader implements rpc.ServerCodec func (j *JSONWebsocketCodec) ReadRequestHeader(r *rpc.Request) error { - err := j.conn.ReadJSON(r) - if err != nil { - close(j.err) + m, err := j.readMessage(nil) + if err == nil { + *r = *m.Request } return err } // ReadRequestBody implements rpc.ServerCodec func (j *JSONWebsocketCodec) ReadRequestBody(v interface{}) error { - err := j.conn.ReadJSON(v) - if err != nil { - close(j.err) - } + _, err := j.readMessage(v) return err } - -// WriteResponse implements rpc.ServerCodec -func (j *JSONWebsocketCodec) WriteResponse(r *rpc.Response, v interface{}) error { - j.Lock() - defer j.Unlock() - - if err := j.conn.WriteJSON(r); err != nil { - return err - } - return j.conn.WriteJSON(v) -} diff --git a/xfer/multi_client.go b/xfer/multi_client.go index d0eb5a21e..d4eceb9fc 100644 --- a/xfer/multi_client.go +++ b/xfer/multi_client.go @@ -3,6 +3,7 @@ package xfer import ( "bytes" "errors" + "fmt" "io" "io/ioutil" "log" @@ -15,13 +16,10 @@ import ( const maxConcurrentGET = 10 // ClientFactory is a thing thats makes AppClients -type ClientFactory func(ProbeConfig, string, string) (AppClient, error) +type ClientFactory func(string, string) (AppClient, error) type multiClient struct { - ProbeConfig - clientFactory ClientFactory - handler ControlHandler mtx sync.Mutex sema semaphore @@ -39,16 +37,16 @@ type clientTuple struct { // AppClient for each one. type MultiAppClient interface { Set(hostname string, endpoints []string) + PipeConnection(appID, pipeID string, pipe Pipe) error + PipeClose(appID, pipeID string) error Stop() Publish(io.Reader) error } // NewMultiAppClient creates a new MultiAppClient. -func NewMultiAppClient(pc ProbeConfig, handler ControlHandler, clientFactory ClientFactory) MultiAppClient { +func NewMultiAppClient(clientFactory ClientFactory) MultiAppClient { return &multiClient{ - ProbeConfig: pc, clientFactory: clientFactory, - handler: handler, sema: newSemaphore(maxConcurrentGET), clients: map[string]AppClient{}, @@ -67,7 +65,7 @@ func (c *multiClient) Set(hostname string, endpoints []string) { c.sema.acquire() defer c.sema.release() - client, err := c.clientFactory(c.ProbeConfig, hostname, endpoint) + client, err := c.clientFactory(hostname, endpoint) if err != nil { log.Printf("Error creating new app client: %v", err) return @@ -96,7 +94,7 @@ func (c *multiClient) Set(hostname string, endpoints []string) { _, ok := c.clients[tuple.ID] if !ok { c.clients[tuple.ID] = tuple.AppClient - tuple.AppClient.ControlConnection(c.handler) + tuple.AppClient.ControlConnection() } } c.ids[hostname] = hostIDs @@ -114,6 +112,29 @@ func (c *multiClient) Set(hostname string, endpoints []string) { } } +func (c *multiClient) withClient(appID string, f func(AppClient) error) error { + c.mtx.Lock() + client, ok := c.clients[appID] + c.mtx.Unlock() + if !ok { + return fmt.Errorf("No such app id: %s", appID) + } + return f(client) +} + +func (c *multiClient) PipeConnection(appID, pipeID string, pipe Pipe) error { + return c.withClient(appID, func(client AppClient) error { + client.PipeConnection(pipeID, pipe) + return nil + }) +} + +func (c *multiClient) PipeClose(appID, pipeID string) error { + return c.withClient(appID, func(client AppClient) error { + return client.PipeClose(pipeID) + }) +} + // Stop the MultiAppClient. func (c *multiClient) Stop() { c.mtx.Lock() diff --git a/xfer/multi_client_test.go b/xfer/multi_client_test.go index aecf8b587..b4e495503 100644 --- a/xfer/multi_client_test.go +++ b/xfer/multi_client_test.go @@ -20,7 +20,7 @@ func (c *mockClient) Details() (xfer.Details, error) { return xfer.Details{ID: c.id}, nil } -func (c *mockClient) ControlConnection(handler xfer.ControlHandler) { +func (c *mockClient) ControlConnection() { c.count++ } @@ -33,12 +33,15 @@ func (c *mockClient) Publish(io.Reader) error { return nil } +func (c *mockClient) PipeConnection(_ string, _ xfer.Pipe) {} +func (c *mockClient) PipeClose(_ string) error { return nil } + var ( a1 = &mockClient{id: "1"} // hostname a, app id 1 a2 = &mockClient{id: "2"} // hostname a, app id 2 b2 = &mockClient{id: "2"} // hostname b, app id 2 (duplicate) b3 = &mockClient{id: "3"} // hostname b, app id 3 - factory = func(_ xfer.ProbeConfig, hostname, target string) (xfer.AppClient, error) { + factory = func(hostname, target string) (xfer.AppClient, error) { switch target { case "a1": return a1, nil @@ -55,9 +58,6 @@ var ( func TestMultiClient(t *testing.T) { var ( - controlHandler = xfer.ControlHandlerFunc(func(_ xfer.Request) xfer.Response { - return xfer.Response{} - }) expect = func(i, j int) { if i != j { _, file, line, _ := runtime.Caller(1) @@ -66,7 +66,7 @@ func TestMultiClient(t *testing.T) { } ) - mp := xfer.NewMultiAppClient(xfer.ProbeConfig{}, controlHandler, factory) + mp := xfer.NewMultiAppClient(factory) defer mp.Stop() // Add two hostnames with overlapping apps, check we don't add the same app twice @@ -88,7 +88,7 @@ func TestMultiClient(t *testing.T) { } func TestMultiClientPublish(t *testing.T) { - mp := xfer.NewMultiAppClient(xfer.ProbeConfig{}, nil, factory) + mp := xfer.NewMultiAppClient(factory) defer mp.Stop() sum := func() int { return a1.publish + a2.publish + b2.publish + b3.publish } diff --git a/xfer/pipes.go b/xfer/pipes.go new file mode 100644 index 000000000..25a061424 --- /dev/null +++ b/xfer/pipes.go @@ -0,0 +1,128 @@ +package xfer + +import ( + "io" + "sync" + + "github.com/gorilla/websocket" +) + +// Pipe is a bi-directional channel from someone thing in the probe +// to the UI. +type Pipe interface { + Ends() (io.ReadWriter, io.ReadWriter) + CopyToWebsocket(io.ReadWriter, *websocket.Conn) error + + Close() error + Closed() bool + OnClose(func()) +} + +type pipe struct { + mtx sync.Mutex + port, starboard io.ReadWriter + quit chan struct{} + closed bool + onClose func() +} + +// NewPipe makes a new... pipe. +func NewPipe() Pipe { + r1, w1 := io.Pipe() + r2, w2 := io.Pipe() + return &pipe{ + port: struct { + io.Reader + io.Writer + }{ + r1, w2, + }, + starboard: struct { + io.Reader + io.Writer + }{ + r2, w1, + }, + quit: make(chan struct{}), + } +} + +func (p *pipe) Ends() (io.ReadWriter, io.ReadWriter) { + return p.port, p.starboard +} + +func (p *pipe) Close() error { + p.mtx.Lock() + var onClose func() + if !p.closed { + p.closed = true + close(p.quit) + onClose = p.onClose + } + p.mtx.Unlock() + + // Don't run onClose under lock. + if onClose != nil { + onClose() + } + return nil +} + +func (p *pipe) Closed() bool { + p.mtx.Lock() + defer p.mtx.Unlock() + return p.closed +} + +func (p *pipe) OnClose(f func()) { + p.mtx.Lock() + defer p.mtx.Unlock() + p.onClose = f +} + +// CopyToWebsocket copies pipe data to/from a websocket. It blocks. +func (p *pipe) CopyToWebsocket(end io.ReadWriter, conn *websocket.Conn) error { + errors := make(chan error, 1) + + // Read-from-UI loop + go func() { + for { + _, buf, err := conn.ReadMessage() // TODO type should be binary message + if err != nil { + errors <- err + return + } + + if _, err := end.Write(buf); err != nil { + errors <- err + return + } + } + }() + + // Write-to-UI loop + go func() { + buf := make([]byte, 1024) + for { + n, err := end.Read(buf) + if err != nil { + errors <- err + return + } + + if err := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil { + errors <- err + return + } + } + }() + + // block until one of the goroutines exits + // this convoluted mechanism is to ensure we only close the websocket once. + select { + case err := <-errors: + return err + case <-p.quit: + return nil + } +}