Add signal handler for ruby socket cleanup on shutdown

2 files changed, 50 insertions, 7 deletions

diff --git a/internal/rubyserver/rubyserver.go b/internal/rubyserver/rubyserver.go
index 6ee92ef10..0c7d349d3 100644
--- a/internal/rubyserver/rubyserver.go
+++ b/internal/rubyserver/rubyserver.go
@@ -30,7 +30,7 @@ const (
 )
 
 var (
-	socketPath string
+	socketDir string
 
 	lazyInit sync.Once
 
@@ -51,21 +51,48 @@ func prepareSocketPath() {
 	// that is not too deep. We need a directory, not a tempfile, because we
 	// will later want to set its permissions to 0700. The permission change
 	// is done in the Ruby child process.
-	socketDir, err := ioutil.TempDir("", "gitaly-ruby")
+	var err error
+	socketDir, err = ioutil.TempDir("", "gitaly-ruby")
 	if err != nil {
 		log.Fatalf("create ruby server socket directory: %v", err)
 	}
-	socketPath = path.Join(filepath.Clean(socketDir), "socket")
+}
+
+func socketPath() string {
+	if socketDir == "" {
+		panic("socketDir is not set")
+	}
+
+	return path.Join(filepath.Clean(socketDir), "socket")
+}
+
+// Server represents a gitaly-ruby helper process.
+type Server struct {
+	*supervisor.Process
+}
+
+// Stop shuts down the gitaly-ruby helper process and cleans up resources.
+func (s *Server) Stop() {
+	if s != nil {
+		if s.Process != nil {
+			s.Process.Stop()
+		}
+	}
+
+	if socketDir != "" {
+		os.RemoveAll(socketDir)
+	}
 }
 
 // Start spawns the Ruby server.
-func Start() (*supervisor.Process, error) {
+func Start() (*Server, error) {
 	lazyInit.Do(prepareSocketPath)
 
-	args := []string{"bundle", "exec", "bin/gitaly-ruby", fmt.Sprintf("%d", os.Getpid()), socketPath}
+	args := []string{"bundle", "exec", "bin/gitaly-ruby", fmt.Sprintf("%d", os.Getpid()), socketPath()}
 	env := append(os.Environ(), "GITALY_RUBY_GIT_BIN_PATH="+helper.GitPath(),
 		fmt.Sprintf("GITALY_RUBY_WRITE_BUFFER_SIZE=%d", streamio.WriteBufferSize))
-	return supervisor.New(env, args, config.Config.Ruby.Dir)
+	p, err := supervisor.New(env, args, config.Config.Ruby.Dir)
+	return &Server{Process: p}, err
 }
 
 // CommitServiceClient returns a CommitServiceClient instance that is
@@ -87,7 +114,7 @@ func DiffServiceClient(ctx context.Context) (pb.DiffServiceClient, error) {
 func newConnection(ctx context.Context) (*grpc.ClientConn, error) {
 	dialCtx, cancel := context.WithTimeout(ctx, ConnectTimeout)
 	defer cancel()
-	return grpc.DialContext(dialCtx, socketPath, dialOptions()...)
+	return grpc.DialContext(dialCtx, socketPath(), dialOptions()...)
 }
 
 func dialOptions() []grpc.DialOption {
diff --git a/internal/rubyserver/rubyserver_test.go b/internal/rubyserver/rubyserver_test.go
new file mode 100644
index 000000000..a638c2ee4
--- /dev/null
+++ b/internal/rubyserver/rubyserver_test.go
@@ -0,0 +1,16 @@
+package rubyserver
+
+import (
+	"testing"
+)
+
+func TestStopSafe(t *testing.T) {
+	badServers := []*Server{
+		nil,
+		&Server{},
+	}
+
+	for _, bs := range badServers {
+		bs.Stop()
+	}
+}