Add daemon mode for wal-push command

cmd/pg/daemon.go

@@ -0,0 +1,55 @@
+package pg
+
+import (
+	"os"
+	"path"
+
+	"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"
+	"github.com/wal-g/wal-g/utility"
+)
+
+const DaemonShortDescription = "Uploads a WAL file to storage"
+
+// daemonCmd represents the daemon archive command
+var daemonCmd = &cobra.Command{
+	Use:   "daemon daemon_socket_path",
+	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)
+		pathToData, exist := os.LookupEnv("PGDATA")
+		if !exist {
+			tracelog.ErrorLogger.Fatal("'PGDATA' variable is not defined or does not exist")
+		}
+		pathToWal := path.Join(pathToData, "pg_wal")
+
+		postgres.HandleDaemon(uploader, args[0], pathToWal)
+	},
+}
+
+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,132 @@
+package postgres
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"net"
+	"os"
+	"path"
+
+	"github.com/wal-g/tracelog"
+)
+
+type CheckType byte
+type FileType byte
+
+type MessageSocketHandler interface {
+	Handle(messageBody []byte, c net.Conn, f func(string) error) error
+}
+
+func MessageTypeConstruct(messageType byte) MessageSocketHandler {
+	switch messageType {
+	case 'C':
+		return CheckType('C')
+	case 'F':
+		return FileType('F')
+	default:
+		return nil
+	}
+}
+
+func (msg CheckType) Handle(messageBody []byte, c net.Conn, f func(string) error) error {
+	_, err := c.Write([]byte{'O', 0, 3})
+	if err != nil {
+		tracelog.ErrorLogger.Printf("Error on writing in socket: %v", err)
+		return err
+	}
+	return nil
+}
+
+func (msg FileType) Handle(messageBody []byte, c net.Conn, f func(string) error) error {
+	if len(messageBody) < 24 {
+		if len(messageBody) > 0 {
+			tracelog.ErrorLogger.Printf("Received incorrect message %s", messageBody)
+		} else {
+			tracelog.ErrorLogger.Println("Received empty message")
+		}
+		return errors.New(fmt.Sprint("Incorrect message accepted"))
+	}
+	tracelog.InfoLogger.Printf("wal file name: %s\n", messageBody)
+	err := f(string(messageBody))
+	if err != nil {
+		tracelog.ErrorLogger.Printf("wal-push failed: %v\n", err)
+		return err
+	}
+	_, err = c.Write([]byte{'O', 0, 3})
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// HandleDaemon is invoked to perform daemon mode
+func HandleDaemon(uploader *WalUploader, pathToSocket string, pathToWal string) {
+	_ = os.Remove(pathToSocket)
+	l, err := net.Listen("unix", pathToSocket)
+	if err != nil {
+		tracelog.ErrorLogger.Fatal("Error on listening socket:", err)
+	}
+	for {
+		fd, err := l.Accept()
+		if err != nil {
+			tracelog.ErrorLogger.Println("Failed to accept, err:", err)
+		}
+		go DaemonProcess(fd, func(walFileName string) error {
+			fullPath := path.Join(pathToWal, walFileName)
+			tracelog.InfoLogger.Printf("starting wal-push for %s\n", fullPath)
+			return HandleWALPush(uploader, fullPath)
+		})
+	}
+}
+
+func DaemonProcess(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)
+	for {
+		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{'E', 0, 3})
+			return
+		}
+
+		byteType, byteLength, byteBody := MessageParser(buf)
+		requestType := MessageTypeConstruct(byteType)
+		err = requestType.Handle(byteBody, c, f)
+		if err != nil {
+			tracelog.ErrorLogger.Println("Failed to handle message, err:", err)
+		}
+		if byteType == 'F' {
+			return
+		}
+		if int(byteLength) < nr {
+			tracelog.InfoLogger.Println("Read remaining buffer...")
+			byteType, byteLength, byteBody = MessageParser(buf[byteLength:])
+			err = requestType.Handle(byteBody, c, f)
+			if err != nil {
+				tracelog.ErrorLogger.Println("Failed to handle message, err:", err)
+			}
+			return
+		}
+	}
+}
+
+// MessageParser is invoked to read bytes from buffer
+func MessageParser(buf []byte) (byte, uint16, []byte) {
+	messageType := buf[0]
+	var messageLength uint16
+	l := bytes.NewReader(buf[1:3])
+	err := binary.Read(l, binary.BigEndian, &messageLength)
+	if err != nil {
+		tracelog.ErrorLogger.Printf("Failed to read message length, err: %v", err)
+	}
+	messageBody := buf[3 : messageLength-1]
+	return messageType, messageLength, messageBody
+}

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