MINOR: runtime: add atomic map replacement

runtime/interface.go

@@ -52,6 +52,8 @@ type Maps interface {
 	// ParseMapEntriesFromFile reads entries from file
 	ParseMapEntriesFromFile(inputFile io.Reader, hasID bool) models.MapEntries
 	AddMapPayloadVersioned(name string, entries models.MapEntries) error
+	// ReplaceMapPayload atomically replaces all entries in a map.
+	ReplaceMapPayload(name string, entries models.MapEntries) error
 	// PrepareMap allocates a new map version
 	PrepareMap(name string) (version string, err error)
 	// CommitMap commits all changes made to a map version

runtime/runtime_client.go

@@ -694,37 +694,27 @@ func (c *client) AddMapPayload(name, payload string) error {
 }
 
 func parseMapPayload(entries models.MapEntries, maxBufSize int) (bool, []string) {
-	prevKV := ""
-	currKV := ""
 	data := ""
 	var payload []string
 	var exceededSize bool
 	for _, d := range entries {
-		if prevKV != "" {
-			data += prevKV
-			prevKV = ""
-		}
 		kv := d.Key + " " + d.Value + "\n"
-		data += kv
-		switch {
-		case len(data) < maxBufSize:
-			currKV = data
-		case len(data) == maxBufSize:
+		if data != "" && len(data)+len(kv) > maxBufSize {
+			exceededSize = true
 			payload = append(payload, data)
 			data = ""
-		case len(data) > maxBufSize:
-			exceededSize = true
-			if currKV == "" {
-				currKV = kv
+		}
+		data += kv
+		if len(data) >= maxBufSize {
+			if len(data) > maxBufSize {
+				exceededSize = true
 			}
-			payload = append(payload, currKV)
-			prevKV = d.Key + " " + d.Value + "\n"
+			payload = append(payload, data)
 			data = ""
-			currKV = ""
 		}
 	}
-	if len(currKV) > 0 {
-		payload = append(payload, currKV)
+	if data != "" {
+		payload = append(payload, data)
 	}
 	return exceededSize, payload
 }
@@ -770,6 +760,38 @@ func (c *client) AddMapPayloadVersioned(name string, entries models.MapEntries)
 	return nil
 }
 
+// ReplaceMapPayload atomically replaces all entries in a map using a prepared map version.
+func (c *client) ReplaceMapPayload(name string, entries models.MapEntries) error {
+	if !c.runtime.IsValid() {
+		return errors.New("no valid runtime found")
+	}
+	v := HAProxyVersion{Version: semver.MustParse("2.4")}
+	if !c.IsVersionBiggerOrEqual(&v) {
+		return fmt.Errorf("not supported for HAProxy versions lower than 2.4 %w", native_errors.ErrGeneral)
+	}
+	name, err := c.GetMapsPath(name)
+	if err != nil {
+		return fmt.Errorf("%s %w", c.runtime.socketPath, err)
+	}
+	_, payload := parseMapPayload(entries, maxBufSize)
+	version, err := c.runtime.PrepareMap(name)
+	if err != nil {
+		return fmt.Errorf("%s %w", c.runtime.socketPath, err)
+	}
+	if err = c.runtime.ClearMapVersioned(name, version); err != nil {
+		return fmt.Errorf("%s %w", c.runtime.socketPath, err)
+	}
+	for i := range payload {
+		if err = c.runtime.AddMapPayloadVersioned(version, name, payload[i]); err != nil {
+			return fmt.Errorf("%s %w", c.runtime.socketPath, err)
+		}
+	}
+	if err = c.runtime.CommitMap(version, name); err != nil {
+		return fmt.Errorf("%s %w", c.runtime.socketPath, err)
+	}
+	return nil
+}
+
 // AddMapEntry adds an entry into the map file
 func (c *client) AddMapEntry(name, key, value string) error {
 	if !c.runtime.IsValid() {

runtime/runtime_client_test.go

@@ -16,12 +16,119 @@
 package runtime
 
 import (
+	"bufio"
+	"errors"
+	"net"
+	"os"
+	"strings"
+	"sync"
 	"testing"
 
+	"github.com/Masterminds/semver/v3"
+	native_errors "github.com/haproxytech/client-native/v6/errors"
 	"github.com/haproxytech/client-native/v6/models"
+	"github.com/haproxytech/client-native/v6/runtime/options"
 	"github.com/stretchr/testify/assert"
 )
 
+type mapCommandSocket struct {
+	net.Listener
+
+	reply    func(string) string
+	mu       sync.Mutex
+	commands []string
+}
+
+func newMapCommandSocket(t *testing.T, reply func(string) string) *mapCommandSocket {
+	t.Helper()
+	f, err := os.CreateTemp("", "client-native-map-*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	addr := f.Name()
+	_ = f.Close()
+	_ = os.Remove(addr)
+	listener, err := net.Listen("unix", addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	socket := &mapCommandSocket{Listener: listener, reply: reply}
+	go socket.serve()
+	t.Cleanup(func() {
+		_ = listener.Close()
+		_ = os.Remove(addr)
+	})
+	return socket
+}
+
+func (s *mapCommandSocket) serve() {
+	for {
+		conn, err := s.Accept()
+		if err != nil {
+			return
+		}
+		go s.handleConnection(conn)
+	}
+}
+
+func (s *mapCommandSocket) handleConnection(conn net.Conn) {
+	defer conn.Close()
+
+	reader := bufio.NewReader(conn)
+	command, err := reader.ReadString('\n')
+	if err != nil {
+		return
+	}
+	if strings.Contains(command, "<<") {
+		for {
+			line, readErr := reader.ReadString('\n')
+			command += line
+			if line == "\n" || readErr != nil {
+				break
+			}
+		}
+	}
+	command = strings.TrimPrefix(command, "set severity-output number;")
+	command = strings.TrimSuffix(command, "\n")
+
+	s.mu.Lock()
+	s.commands = append(s.commands, command)
+	s.mu.Unlock()
+
+	response := ""
+	if s.reply != nil {
+		response = s.reply(command)
+	}
+	if response == "" && strings.HasPrefix(command, "prepare map ") {
+		response = "map transaction: version: 1"
+	}
+	_, _ = conn.Write([]byte(response))
+}
+
+func (s *mapCommandSocket) Commands() []string {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return append([]string(nil), s.commands...)
+}
+
+func newMapClient(t *testing.T, socket *mapCommandSocket, version string) *client {
+	t.Helper()
+	runtime := &SingleRuntime{}
+	if err := runtime.Init(socket.Addr().String(), false, options.RuntimeOptions{DoNotCheckRuntimeOnInit: true}); err != nil {
+		t.Fatal(err)
+	}
+	mapsDir := "/etc/haproxy/maps"
+	previousVersion := haproxyVersion
+	haproxyVersion = &HAProxyVersion{Version: semver.MustParse(version)}
+	t.Cleanup(func() {
+		haproxyVersion = previousVersion
+	})
+	return &client{
+		options: options.RuntimeOptions{MapsDir: &mapsDir},
+		runtime: runtime,
+	}
+}
+
 func Test_parseMapPayload(t *testing.T) {
 	type args struct {
 		entries    models.MapEntries
@@ -65,6 +172,43 @@ func Test_parseMapPayload(t *testing.T) {
 				"key2 val2\n",
 			},
 		},
+		{
+			name: "Keep final entry after crossing maxBufSize",
+			args: args{entries: models.MapEntries{
+				&models.MapEntry{Key: "a", Value: "bbbb"},
+				&models.MapEntry{Key: "b", Value: "bbbb"},
+			}, maxBufSize: 10},
+			wantExceededSize: true,
+			wantPayload: []string{
+				"a bbbb\n",
+				"b bbbb\n",
+			},
+		},
+		{
+			name: "Keep entries across repeated maxBufSize crossings",
+			args: args{entries: models.MapEntries{
+				&models.MapEntry{Key: "a", Value: "bbbb"},
+				&models.MapEntry{Key: "b", Value: "bbbb"},
+				&models.MapEntry{Key: "c", Value: "bbbb"},
+			}, maxBufSize: 10},
+			wantExceededSize: true,
+			wantPayload: []string{
+				"a bbbb\n",
+				"b bbbb\n",
+				"c bbbb\n",
+			},
+		},
+		{
+			name: "Do not repeat data after reaching maxBufSize exactly",
+			args: args{entries: models.MapEntries{
+				&models.MapEntry{Key: "a", Value: "b"},
+				&models.MapEntry{Key: "c", Value: "ddd"},
+			}, maxBufSize: 10},
+			wantExceededSize: false,
+			wantPayload: []string{
+				"a b\nc ddd\n",
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -78,3 +222,99 @@ func Test_parseMapPayload(t *testing.T) {
 		})
 	}
 }
+
+func TestReplaceMapPayload(t *testing.T) {
+	socket := newMapCommandSocket(t, nil)
+	client := newMapClient(t, socket, "2.4.0")
+	entries := models.MapEntries{
+		&models.MapEntry{Key: "first", Value: "one"},
+		&models.MapEntry{Key: "second", Value: "two"},
+	}
+
+	if err := client.ReplaceMapPayload("example", entries); err != nil {
+		t.Fatalf("ReplaceMapPayload() error = %v", err)
+	}
+
+	path := "/etc/haproxy/maps/example.map"
+	assert.Equal(t, []string{
+		"prepare map " + path,
+		"clear map @1 " + path,
+		"add map @1 " + path + " <<\nfirst one\nsecond two\n",
+		"commit map @1 " + path,
+	}, socket.Commands())
+}
+
+func TestReplaceMapPayloadMultipleChunksPreservesOrder(t *testing.T) {
+	socket := newMapCommandSocket(t, nil)
+	client := newMapClient(t, socket, "2.4.0")
+	firstValue := strings.Repeat("a", maxBufSize-len("first")-2)
+	secondValue := strings.Repeat("b", maxBufSize-len("second")-2)
+	entries := models.MapEntries{
+		&models.MapEntry{Key: "first", Value: firstValue},
+		&models.MapEntry{Key: "second", Value: secondValue},
+	}
+
+	if err := client.ReplaceMapPayload("example", entries); err != nil {
+		t.Fatalf("ReplaceMapPayload() error = %v", err)
+	}
+
+	path := "/etc/haproxy/maps/example.map"
+	assert.Equal(t, []string{
+		"prepare map " + path,
+		"clear map @1 " + path,
+		"add map @1 " + path + " <<\nfirst " + firstValue + "\n",
+		"add map @1 " + path + " <<\nsecond " + secondValue + "\n",
+		"commit map @1 " + path,
+	}, socket.Commands())
+}
+
+func TestReplaceMapPayloadEmpty(t *testing.T) {
+	socket := newMapCommandSocket(t, nil)
+	client := newMapClient(t, socket, "2.4.0")
+
+	if err := client.ReplaceMapPayload("example", nil); err != nil {
+		t.Fatalf("ReplaceMapPayload() error = %v", err)
+	}
+
+	path := "/etc/haproxy/maps/example.map"
+	assert.Equal(t, []string{
+		"prepare map " + path,
+		"clear map @1 " + path,
+		"commit map @1 " + path,
+	}, socket.Commands())
+}
+
+func TestReplaceMapPayloadUnsupported(t *testing.T) {
+	socket := newMapCommandSocket(t, nil)
+	client := newMapClient(t, socket, "2.3.0")
+
+	err := client.ReplaceMapPayload("example", nil)
+	if !errors.Is(err, native_errors.ErrGeneral) {
+		t.Fatalf("ReplaceMapPayload() error = %v, want %v", err, native_errors.ErrGeneral)
+	}
+	assert.Empty(t, socket.Commands())
+}
+
+func TestReplaceMapPayloadDoesNotCommitAfterAddFailure(t *testing.T) {
+	socket := newMapCommandSocket(t, func(command string) string {
+		if strings.HasPrefix(command, "add map ") {
+			return "[3]: add failed"
+		}
+		return ""
+	})
+	client := newMapClient(t, socket, "2.4.0")
+
+	err := client.ReplaceMapPayload("example", models.MapEntries{
+		&models.MapEntry{Key: "first", Value: "one"},
+	})
+	if !errors.Is(err, native_errors.ErrGeneral) {
+		t.Fatalf("ReplaceMapPayload() error = %v, want %v", err, native_errors.ErrGeneral)
+	}
+
+	path := "/etc/haproxy/maps/example.map"
+	assert.Equal(t, []string{
+		"prepare map " + path,
+		"clear map @1 " + path,
+		"add map @1 " + path + " <<\nfirst one\n",
+	}, socket.Commands())
+}