Add daemon mode for wal-push command

cmd/pg/daemon.go

@@ -0,0 +1,70 @@
+package pg
+
+import (
+	"net"
+	"os"
+	"path"
+
+	"github.com/wal-g/wal-g/utility"
+
+	"github.com/spf13/cobra"
+	"github.com/wal-g/tracelog"
+	"github.com/wal-g/wal-g/internal"
+	"github.com/wal-g/wal-g/internal/asm"
+	"github.com/wal-g/wal-g/internal/databases/postgres"
+)
+
+const DaemonShortDescription = "Uploads a WAL file to storage"
+const DaemonSocketName = "/tmp/wal-push.sock"
+
+// daemonCmd represents the daemon command
+var daemonCmd = &cobra.Command{
+	Use:   "daemon wal_filepath",
+	Short: DaemonShortDescription, // TODO : improve description
+	Args:  cobra.ExactArgs(1),
+	Run: func(cmd *cobra.Command, args []string) {
+		uploader, err := postgres.ConfigureWalUploader()
+		tracelog.ErrorLogger.FatalOnError(err)
+
+		archiveStatusManager, err := internal.ConfigureArchiveStatusManager()
+		if err == nil {
+			uploader.ArchiveStatusManager = asm.NewDataFolderASM(archiveStatusManager)
+		} else {
+			tracelog.ErrorLogger.PrintError(err)
+			uploader.ArchiveStatusManager = asm.NewNopASM()
+		}
+
+		PGArchiveStatusManager, err := internal.ConfigurePGArchiveStatusManager()
+		if err == nil {
+			uploader.PGArchiveStatusManager = asm.NewDataFolderASM(PGArchiveStatusManager)
+		} else {
+			tracelog.ErrorLogger.PrintError(err)
+			uploader.PGArchiveStatusManager = asm.NewNopASM()
+		}
+
+		uploader.UploadingFolder = uploader.UploadingFolder.GetSubFolder(utility.WalPath)
+
+		_ = os.Remove(DaemonSocketName)
+		l, err := net.Listen("unix", DaemonSocketName)
+		if err != nil {
+			tracelog.ErrorLogger.Fatal("listen error:", err)
+		}
+		for {
+			fd, err := l.Accept()
+			if err != nil {
+				tracelog.ErrorLogger.Println("accept error:", err)
+			}
+
+			go postgres.HandleDaemon(fd, func(walFileName string) error {
+				fullPath := path.Join(args[0], walFileName)
+				tracelog.InfoLogger.Printf("starting wal-push for %s\n", fullPath)
+				uploader.UploadingFolder = uploader.UploadingFolder.GetSubFolder(utility.WalPath)
+				return postgres.HandleWALPush(uploader, fullPath)
+			})
+		}
+	},
+}
+
+func init() {
+	Cmd.AddCommand(daemonCmd)
+}

cmd/pg/wal_push.go

@@ -6,6 +6,7 @@ import (
 	"github.com/wal-g/wal-g/internal"
 	"github.com/wal-g/wal-g/internal/asm"
 	"github.com/wal-g/wal-g/internal/databases/postgres"
+	"github.com/wal-g/wal-g/utility"
 )
 
 const WalPushShortDescription = "Uploads a WAL file to storage"
@@ -35,7 +36,11 @@ var walPushCmd = &cobra.Command{
 			uploader.PGArchiveStatusManager = asm.NewNopASM()
 		}
 
-		postgres.HandleWALPush(uploader, args[0])
+		uploader.UploadingFolder = uploader.UploadingFolder.GetSubFolder(utility.WalPath)
+		err = postgres.HandleWALPush(uploader, args[0])
+		if err != nil {
+			tracelog.ErrorLogger.FatalOnError(err)
+		}
 	},
 }
 

docker/pg_tests/scripts/tests/daemon_test.sh

@@ -0,0 +1,28 @@
+#!/bin/sh
+set -e -x
+
+CONFIG_FILE="/tmp/configs/wal_perftest_config.json"
+
+COMMON_CONFIG="/tmp/configs/common_config.json"
+TMP_CONFIG="/tmp/configs/tmp_config.json"
+cat ${CONFIG_FILE} > ${TMP_CONFIG}
+
+echo "," >> ${TMP_CONFIG}
+cat ${COMMON_CONFIG} >> ${TMP_CONFIG}
+/tmp/scripts/wrap_config_file.sh ${TMP_CONFIG}
+
+WAL=$(ls -l ${PGDATA}/pg_wal | head -n2 | tail -n1 | egrep -o "[0-9A-F]{24}")
+
+wal-g --config=${TMP_CONFIG} daemon "${PGDATA}"/pg_wal
+
+# shellcheck disable=SC2039
+nc -U /tmp/wal-push.sock <<< 'CHECK'
+sleep 1
+
+# shellcheck disable=SC2039
+nc -U /tmp/wal-push.sock <<< "${WAL}"
+sleep 1
+
+wal-g st ls WALG_S3_PREFIX/wal_005/"${WAL}".br
+
+

internal/databases/postgres/daemon_handler.go

@@ -0,0 +1,62 @@
+package postgres
+
+import (
+	"net"
+
+	"github.com/wal-g/tracelog"
+)
+
+// HandleDaemon is invoked to perform daemon mode
+func HandleDaemon(c net.Conn, f func(string) error) {
+	defer func(c net.Conn) {
+		err := c.Close()
+		if err != nil {
+			tracelog.ErrorLogger.Printf("Failed to close connection with %s, err: %v\n", c.RemoteAddr(), err)
+		}
+	}(c)
+	buf := make([]byte, 512)
+	nr, err := c.Read(buf)
+	if err != nil {
+		tracelog.ErrorLogger.Printf("Failed to read checking message from client %s, err: %v\n", c.RemoteAddr(), err)
+		_, _ = c.Write([]byte("READ_FAILED"))
+		return
+	}
+	if nr == 5 && string(buf[0:5]) == "CHECK" {
+		_, _ = c.Write([]byte("CHECKED"))
+		tracelog.InfoLogger.Printf("Successful configuration check")
+	} else {
+		tracelog.ErrorLogger.Printf("Error on configuration check")
+		return
+	}
+	n, err := c.Read(buf)
+	if err != nil {
+		tracelog.ErrorLogger.Printf("Failed to read message with file from client %s, err: %v\n", c.RemoteAddr(), err)
+		_, _ = c.Write([]byte("READ_FAILED"))
+		return
+	}
+
+	if n < 24 {
+		if n > 0 {
+			tracelog.ErrorLogger.Printf("Received incorrect message %s from %s", buf[0:n], c.RemoteAddr())
+		} else {
+			tracelog.ErrorLogger.Printf("Received empty message from %s", c.RemoteAddr())
+		}
+		_, _ = c.Write([]byte("BAD_MSG"))
+		return
+	}
+
+	data := buf[0:24]
+	tracelog.InfoLogger.Printf("wal file name: %s\n", string(data))
+
+	err = f(string(data))
+	if err != nil {
+		tracelog.ErrorLogger.Printf("wal-push failed: %v\n", err)
+		_, _ = c.Write([]byte("FAIL"))
+		return
+	}
+
+	_, err = c.Write([]byte("OK"))
+	if err != nil {
+		tracelog.ErrorLogger.Println("OK write fail: ", err)
+	}
+}

internal/databases/postgres/wal_push_handler.go

@@ -13,7 +13,6 @@ import (
 	"github.com/pkg/errors"
 	"github.com/spf13/viper"
 	"github.com/wal-g/tracelog"
-	"github.com/wal-g/wal-g/utility"
 )
 
 type CantOverwriteWalFileError struct {
@@ -32,21 +31,20 @@ func (err CantOverwriteWalFileError) Error() string {
 
 // TODO : unit tests
 // HandleWALPush is invoked to perform wal-g wal-push
-func HandleWALPush(uploader *WalUploader, walFilePath string) {
-	uploader.UploadingFolder = uploader.UploadingFolder.GetSubFolder(utility.WalPath)
+func HandleWALPush(uploader *WalUploader, walFilePath string) error {
 	if uploader.ArchiveStatusManager.IsWalAlreadyUploaded(walFilePath) {
 		err := uploader.ArchiveStatusManager.UnmarkWalFile(walFilePath)
 
 		if err != nil {
 			tracelog.ErrorLogger.Printf("unmark wal-g status for %s file failed due following error %+v", walFilePath, err)
 		}
-		err = uploadLocalWalMetadata(walFilePath, uploader.Uploader)
-		tracelog.ErrorLogger.FatalOnError(err)
-		return
+		return uploadLocalWalMetadata(walFilePath, uploader.Uploader)
 	}
 
 	concurrency, err := internal.GetMaxUploadConcurrency()
-	tracelog.ErrorLogger.FatalOnError(err)
+	if err != nil {
+		return err
+	}
 
 	totalBgUploadedLimit := viper.GetInt32(internal.TotalBgUploadedLimit)
 	// .history files must not be overwritten, see https://github.com/wal-g/wal-g/issues/420
@@ -58,16 +56,23 @@ func HandleWALPush(uploader *WalUploader, walFilePath string) {
 	bgUploader.Start()
 
 	err = uploadWALFile(uploader, walFilePath, bgUploader.preventWalOverwrite)
-	tracelog.ErrorLogger.FatalOnError(err)
+	if err != nil {
+		return err
+	}
 	err = uploadLocalWalMetadata(walFilePath, uploader.Uploader)
-	tracelog.ErrorLogger.FatalOnError(err)
+	if err != nil {
+		return err
+	}
 
 	err = bgUploader.Stop()
-	tracelog.ErrorLogger.FatalOnError(err)
+	if err != nil {
+		return err
+	}
 
 	if uploader.getUseWalDelta() {
 		uploader.FlushFiles()
 	}
+	return nil
 }
 
 // TODO : unit tests