diff --git a/app/multitenant/consul_client.go b/app/multitenant/consul_client.go index 222b5f5d4..a0a47676f 100644 --- a/app/multitenant/consul_client.go +++ b/app/multitenant/consul_client.go @@ -20,6 +20,7 @@ type ConsulClient interface { Get(key string, out interface{}) error CAS(key string, out interface{}, f CASCallback) error WatchPrefix(prefix string, out interface{}, done chan struct{}, f func(string, interface{}) bool) + DeleteSelected(prefix string, out interface{}, f func(string, interface{}) bool) error } // CASCallback is the type of the callback to CAS. If err is nil, out must be non-nil. @@ -51,6 +52,7 @@ type kv interface { CAS(p *consul.KVPair, q *consul.WriteOptions) (bool, *consul.WriteMeta, error) Get(key string, q *consul.QueryOptions) (*consul.KVPair, *consul.QueryMeta, error) List(prefix string, q *consul.QueryOptions) (consul.KVPairs, *consul.QueryMeta, error) + Delete(key string, w *consul.WriteOptions) (*consul.WriteMeta, error) } type consulClient struct { @@ -179,3 +181,22 @@ func (c *consulClient) WatchPrefix(prefix string, out interface{}, done chan str } } } + +func (c *consulClient) DeleteSelected(prefix string, out interface{}, f func(string, interface{}) bool) error { + kvps, _, err := c.kv.List(prefix, nil) + if err != nil { + return err + } + for _, kvp := range kvps { + if err := json.NewDecoder(bytes.NewReader(kvp.Value)).Decode(out); err != nil { + return fmt.Errorf("Error deserialising %s: %v", kvp.Key, err) + } + if f(kvp.Key, out) { + _, err := c.kv.Delete(kvp.Key, nil) + if err != nil { + return fmt.Errorf("Error deleting %s: %v", kvp.Key, err) + } + } + } + return nil +} diff --git a/app/multitenant/consul_pipe_router.go b/app/multitenant/consul_pipe_router.go index 4c01b104a..ec53a82d8 100644 --- a/app/multitenant/consul_pipe_router.go +++ b/app/multitenant/consul_pipe_router.go @@ -103,13 +103,16 @@ func NewConsulPipeRouter(client ConsulClient, prefix, advertise string, userIDer quit: make(chan struct{}), } - pipeRouter.wait.Add(2) - go pipeRouter.watchAll() - go pipeRouter.actor() - go pipeRouter.privateAPI() return pipeRouter } +func (pr *consulPipeRouter) Start() { + pr.wait.Add(2) + go pr.watchAll() + go pr.actor() + go pr.privateAPI() +} + func (pr *consulPipeRouter) Stop() { close(pr.quit) pr.wait.Wait() @@ -364,6 +367,15 @@ func (pr *consulPipeRouter) Delete(ctx context.Context, id string) error { }) } +// Remove any pipe entries in Consul that were deleted a while ago. +func (pr *consulPipeRouter) CleanUp() error { + threshold := time.Now().Add(-gcTimeout) + return pr.client.DeleteSelected(pr.prefix, &consulPipe{}, func(key string, value interface{}) bool { + cp := *value.(*consulPipe) + return !cp.DeletedAt.IsZero() && cp.DeletedAt.Before(threshold) + }) +} + // A bridgeConnection represents a connection between two pipe router replicas. // They are created & destroyed in response to events from consul, which in turn // are triggered when UIs or Probes connect to various pipe routers. diff --git a/app/multitenant/consul_pipe_router_internal_test.go b/app/multitenant/consul_pipe_router_internal_test.go index 32e05466d..d55e657d0 100644 --- a/app/multitenant/consul_pipe_router_internal_test.go +++ b/app/multitenant/consul_pipe_router_internal_test.go @@ -175,6 +175,7 @@ func TestPipeRouter(t *testing.T) { for i := 0; i < replicas; i++ { pr := NewConsulPipeRouter(consul, "", fmt.Sprintf("127.0.0.1:44%02d", i), NoopUserIDer) + pr.Start() defer pr.Stop() pt.prs = append(pt.prs, pr) } diff --git a/app/multitenant/mock_consul_client_internal_test.go b/app/multitenant/mock_consul_client_internal_test.go index 237f76aa0..9052ea21c 100644 --- a/app/multitenant/mock_consul_client_internal_test.go +++ b/app/multitenant/mock_consul_client_internal_test.go @@ -65,6 +65,13 @@ func (m *mockKV) CAS(p *consul.KVPair, q *consul.WriteOptions) (bool, *consul.Wr return true, nil, nil } +func (m *mockKV) Delete(key string, q *consul.WriteOptions) (*consul.WriteMeta, error) { + m.mtx.Lock() + delete(m.kvps, key) + m.mtx.Unlock() + return nil, nil +} + func (m *mockKV) Get(key string, q *consul.QueryOptions) (*consul.KVPair, *consul.QueryMeta, error) { m.mtx.Lock() defer m.mtx.Unlock() diff --git a/app/pipe_router.go b/app/pipe_router.go index 273593175..118cd0fb0 100644 --- a/app/pipe_router.go +++ b/app/pipe_router.go @@ -41,6 +41,8 @@ type PipeRouter interface { Get(context.Context, string, End) (xfer.Pipe, io.ReadWriter, error) Release(context.Context, string, End) error Delete(context.Context, string) error + CleanUp() error + Start() Stop() } @@ -81,11 +83,14 @@ func NewLocalPipeRouter() PipeRouter { quit: make(chan struct{}), pipes: map[string]*pipe{}, } - pipeRouter.wait.Add(1) - go pipeRouter.gcLoop() return pipeRouter } +func (pr *localPipeRouter) Start() { + pr.wait.Add(1) + go pr.gcLoop() +} + func (pr *localPipeRouter) Exists(_ context.Context, id string) (bool, error) { pr.Lock() defer pr.Unlock() @@ -151,6 +156,10 @@ func (pr *localPipeRouter) Delete(_ context.Context, id string) error { return nil } +func (pr *localPipeRouter) CleanUp() error { + return nil +} + func (pr *localPipeRouter) Stop() { close(pr.quit) pr.wait.Wait() diff --git a/app/pipes_internal_test.go b/app/pipes_internal_test.go index da74b8aac..5d465e555 100644 --- a/app/pipes_internal_test.go +++ b/app/pipes_internal_test.go @@ -26,7 +26,6 @@ func TestPipeTimeout(t *testing.T) { router := mux.NewRouter() pr := NewLocalPipeRouter().(*localPipeRouter) RegisterPipeRoutes(router, pr) - pr.Stop() // we don't want the loop running in the background mtime.NowForce(time.Now()) defer mtime.NowReset() @@ -71,6 +70,7 @@ func TestPipeClose(t *testing.T) { router := mux.NewRouter() pr := NewLocalPipeRouter() RegisterPipeRoutes(router, pr) + pr.Start() defer pr.Stop() server := httptest.NewServer(router) diff --git a/prog/app.go b/prog/app.go index 31c7847ed..2a96e7bb7 100644 --- a/prog/app.go +++ b/prog/app.go @@ -267,6 +267,12 @@ func appMain(flags appFlags) { log.Fatalf("Error creating pipe router: %v", err) return } + if flags.cleanupPipesOnly { + pipeRouter.CleanUp() + return + } + + pipeRouter.Start() // Start background version checking checkpoint.CheckInterval(&checkpoint.CheckParams{ diff --git a/prog/main.go b/prog/main.go index e44755228..50ddb9918 100644 --- a/prog/main.go +++ b/prog/main.go @@ -168,6 +168,7 @@ type appFlags struct { controlRouterURL string controlRPCTimeout time.Duration pipeRouterURL string + cleanupPipesOnly bool natsHostname string memcachedHostname string memcachedTimeout time.Duration @@ -381,6 +382,7 @@ func setupFlags(flags *flags) { flag.StringVar(&flags.app.controlRouterURL, "app.control.router", "local", "Control router to use (local or sqs)") flag.DurationVar(&flags.app.controlRPCTimeout, "app.control.rpctimeout", time.Minute, "Timeout for control RPC") flag.StringVar(&flags.app.pipeRouterURL, "app.pipe.router", "local", "Pipe router to use (local)") + flag.BoolVar(&flags.app.cleanupPipesOnly, "app.pipe.cleanup-only", false, "Clean up deleted pipes and exit") flag.StringVar(&flags.app.natsHostname, "app.nats", "", "Hostname for NATS service to use for shortcut reports. If empty, shortcut reporting will be disabled.") flag.StringVar(&flags.app.memcachedHostname, "app.memcached.hostname", "", "Hostname for memcached service to use when caching reports. If empty, no memcached will be used.") flag.DurationVar(&flags.app.memcachedTimeout, "app.memcached.timeout", 100*time.Millisecond, "Maximum time to wait before giving up on memcached requests.")