From 883afeaee39bb6fb2cc519a023f14bbafeb7f16a Mon Sep 17 00:00:00 2001
From: Diego Ciangottini <diego.ciangottini@pg.infn.it>
Date: Sun, 22 Sep 2024 09:20:04 +0200
Subject: [PATCH] listen on socket

Signed-off-by: Diego Ciangottini <diego.ciangottini@pg.infn.it>
---
 cmd/main.go                      | 36 +++++++++++++++++++++++++++++---
 docker/SlurmConfig.yaml          |  3 ++-
 examples/config/SlurmConfig.yaml |  3 ++-
 pkg/slurm/types.go               |  1 +
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index c21d421..3276f27 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -4,9 +4,13 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"strconv"
+	"strings"
+	"syscall"
 	"time"
 
 	"github.com/sirupsen/logrus"
@@ -135,8 +139,34 @@ func main() {
 	SidecarAPIs.CreateDirectories()
 	SidecarAPIs.LoadJIDs()
 
-	err = http.ListenAndServe(":"+slurmConfig.Sidecarport, mutex)
-	if err != nil {
-		log.G(ctx).Fatal(err)
+	if strings.HasPrefix(slurmConfig.Socket, "unix://") {
+		// Create a Unix domain socket and listen for incoming connections.
+		socket, err := net.Listen("unix", strings.ReplaceAll(slurmConfig.Socket, "unix://", ""))
+		if err != nil {
+			panic(err)
+		}
+
+		// Cleanup the sockfile.
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+		go func() {
+			<-c
+			os.Remove(strings.ReplaceAll(slurmConfig.Socket, "unix://", ""))
+			os.Exit(1)
+		}()
+		server := http.Server{
+			Handler: mutex,
+		}
+
+		log.G(ctx).Info(socket)
+
+		if err := server.Serve(socket); err != nil {
+			log.G(ctx).Fatal(err)
+		}
+	} else {
+		err = http.ListenAndServe(":"+slurmConfig.Sidecarport, mutex)
+		if err != nil {
+			log.G(ctx).Fatal(err)
+		}
 	}
 }
diff --git a/docker/SlurmConfig.yaml b/docker/SlurmConfig.yaml
index 1879e25..46e86d0 100644
--- a/docker/SlurmConfig.yaml
+++ b/docker/SlurmConfig.yaml
@@ -1,4 +1,5 @@
 SidecarPort: "4000"
+Socket: ""
 SbatchPath: "/usr/bin/sbatch"
 ScancelPath: "/usr/bin/scancel"
 SqueuePath: "/usr/bin/squeue"
@@ -12,4 +13,4 @@ TsocksPath: "$WORK/tsocks-1.8beta5+ds1/libtsocks.so"
 TsocksLoginNode: "login01"
 BashPath: /bin/bash
 VerboseLogging: true
-ErrorsOnlyLogging: false
\ No newline at end of file
+ErrorsOnlyLogging: false
diff --git a/examples/config/SlurmConfig.yaml b/examples/config/SlurmConfig.yaml
index 1879e25..46e86d0 100644
--- a/examples/config/SlurmConfig.yaml
+++ b/examples/config/SlurmConfig.yaml
@@ -1,4 +1,5 @@
 SidecarPort: "4000"
+Socket: ""
 SbatchPath: "/usr/bin/sbatch"
 ScancelPath: "/usr/bin/scancel"
 SqueuePath: "/usr/bin/squeue"
@@ -12,4 +13,4 @@ TsocksPath: "$WORK/tsocks-1.8beta5+ds1/libtsocks.so"
 TsocksLoginNode: "login01"
 BashPath: /bin/bash
 VerboseLogging: true
-ErrorsOnlyLogging: false
\ No newline at end of file
+ErrorsOnlyLogging: false
diff --git a/pkg/slurm/types.go b/pkg/slurm/types.go
index 8482dd3..d0f14e5 100644
--- a/pkg/slurm/types.go
+++ b/pkg/slurm/types.go
@@ -7,6 +7,7 @@ type SlurmConfig struct {
 	Scancelpath       string `yaml:"ScancelPath"`
 	Squeuepath        string `yaml:"SqueuePath"`
 	Sidecarport       string `yaml:"SidecarPort"`
+	Socket            string `yaml:"Socket"`
 	ExportPodData     bool   `yaml:"ExportPodData"`
 	Commandprefix     string `yaml:"CommandPrefix"`
 	DataRootFolder    string `yaml:"DataRootFolder"`