support unix socket in coordinator

pg-sharding · Apr 12, 2024 · 7f59bb8 · 7f59bb8
1 parent 3301804
commit 7f59bb8
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 7 deletions.
diff --git a/coordinator/app/app.go b/coordinator/app/app.go
@@ -2,10 +2,14 @@ package app
 
 import (
 	"context"
+	"fmt"
 	"net"
+	"os"
+	"path"
 	"sync"
 
 	"github.com/pg-sharding/spqr/pkg/spqrlog"
+	"github.com/pg-sharding/spqr/router/port"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
@@ -55,6 +59,12 @@ func (app *App) Run(withPsql bool) error {
 			}
 		}(wg)
 	}
+	wg.Add(1)
+	go func(wg *sync.WaitGroup) {
+		if err := app.ServeUnixSocket(wg); err != nil {
+			spqrlog.Zero.Error().Err(err).Msg("")
+		}
+	}(wg)
 
 	wg.Wait()
 
@@ -104,7 +114,7 @@ func (app *App) ServeCoordinator(wg *sync.WaitGroup) error {
 				go func() {
 					defer app.sem.Release(1)
 
-					err := app.coordinator.ProcClient(context.TODO(), conn)
+					err := app.coordinator.ProcClient(context.TODO(), conn, port.DefaultRouterPortType)
 					if err != nil {
 						spqrlog.Zero.Error().Err(err).Msg("failed to serve client")
 					}
@@ -150,3 +160,39 @@ func (app *App) ServeGrpcApi(wg *sync.WaitGroup) error {
 
 	return serv.Serve(listener)
 }
+
+func (app *App) ServeUnixSocket(wg *sync.WaitGroup) error {
+	defer wg.Done()
+
+	if err := os.MkdirAll(config.UnixSocketDirectory, 0777); err != nil {
+		return err
+	}
+	socketPath := path.Join(config.UnixSocketDirectory, fmt.Sprintf(".s.PGSQL.%s", config.CoordinatorConfig().CoordinatorPort))
+	lAddr := &net.UnixAddr{Name: socketPath, Net: "unix"}
+	listener, err := net.ListenUnix("unix", lAddr)
+	if err != nil {
+		return err
+	}
+
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			spqrlog.Zero.Error().Err(err).Msg("")
+			continue
+		}
+
+		if err := app.sem.Acquire(context.Background(), 1); err != nil {
+			spqrlog.Zero.Error().Err(err).Msg("")
+			continue
+		}
+
+		go func() {
+			defer app.sem.Release(1)
+
+			err := app.coordinator.ProcClient(context.TODO(), conn, port.UnixSocketPortType)
+			if err != nil {
+				spqrlog.Zero.Error().Err(err).Msg("failed to serve client")
+			}
+		}()
+	}
+}
diff --git a/coordinator/provider/coordinator.go b/coordinator/provider/coordinator.go
@@ -1051,10 +1051,14 @@ func (qc *qdbCoordinator) RemoveTaskGroup(ctx context.Context) error {
 }
 
 // TODO : unit tests
-func (qc *qdbCoordinator) PrepareClient(nconn net.Conn) (CoordinatorClient, error) {
-	cl := psqlclient.NewPsqlClient(nconn, port.DefaultRouterPortType, "")
+func (qc *qdbCoordinator) PrepareClient(nconn net.Conn, pt port.RouterPortType) (CoordinatorClient, error) {
+	cl := psqlclient.NewPsqlClient(nconn, pt, "")
 
-	if err := cl.Init(qc.tlsconfig); err != nil {
+	tlsconfig := qc.tlsconfig
+	if pt == port.UnixSocketPortType {
+		tlsconfig = nil
+	}
+	if err := cl.Init(tlsconfig); err != nil {
 		return nil, err
 	}
 
@@ -1094,8 +1098,8 @@ func (qc *qdbCoordinator) PrepareClient(nconn net.Conn) (CoordinatorClient, erro
 }
 
 // TODO : unit tests
-func (qc *qdbCoordinator) ProcClient(ctx context.Context, nconn net.Conn) error {
-	cl, err := qc.PrepareClient(nconn)
+func (qc *qdbCoordinator) ProcClient(ctx context.Context, nconn net.Conn, pt port.RouterPortType) error {
+	cl, err := qc.PrepareClient(nconn, pt)
 	if err != nil {
 		spqrlog.Zero.Error().Err(err).Msg("")
 		return err

diff --git a/docker/coordinator/Dockerfile b/docker/coordinator/Dockerfile
@@ -1,3 +1,5 @@
 FROM spqr-base-image
 
+RUN apt-get update && apt-get install -y postgresql-client
+
 ENTRYPOINT /spqr/spqr-coordinator -c ${COORDINATOR_CONFIG=/spqr/docker/coordinator/cfg.yaml}
diff --git a/pkg/clientinteractor/interactor.go b/pkg/clientinteractor/interactor.go
@@ -18,6 +18,7 @@ import (
 	"github.com/pg-sharding/spqr/pkg/pool"
 	"github.com/pg-sharding/spqr/pkg/shard"
 	"github.com/pg-sharding/spqr/pkg/txstatus"
+	"github.com/pg-sharding/spqr/router/port"
 	"github.com/pg-sharding/spqr/router/statistics"
 	spqrparser "github.com/pg-sharding/spqr/yacc/console"
 
@@ -31,7 +32,7 @@ import (
 )
 
 type Interactor interface {
-	ProcClient(ctx context.Context, nconn net.Conn) error
+	ProcClient(ctx context.Context, nconn net.Conn, pt port.RouterPortType) error
 }
 
 type PSQLInteractor struct {