1
package server
2
 
3
import (
4
	"fmt"
5
	"io"
6
	"os"
7
	"path/filepath"
8
	"strings"
9
)
10
 
11
const kvProviderMarker = "kv-provider"
12
 
13
// preflightKVProvider establishes identity for the entire server data directory
14
// before any provider opens files or any vnode joins the ring. Layout knowledge
15
// belongs here, where provider selection happens, rather than in the providers.
16
func preflightKVProvider(option, dir string) error {
17
	switch option {
18
	case "memory", "aof", "sqlite":
19
	default:
20
		return fmt.Errorf("unknown kv provider: %s", option)
21
	}
22
	if dir == "" {
23
		return fmt.Errorf("empty data directory")
24
	}
25
 
26
	marker, err := readKVProviderMarker(dir)
27
	if err != nil {
28
		return err
29
	}
30
	layout, err := inspectKVLayout(dir)
31
	if err != nil {
32
		return err
33
	}
34
	if marker != "" && layout != "" && marker != layout {
35
		return fmt.Errorf("storage provider marker identifies %s but data directory %q contains %s storage; preserve the directory and resolve the conflicting stores through recovery before restarting", marker, dir, layout)
36
	}
37
	for _, existing := range []string{marker, layout} {
38
		if existing != "" && existing != option {
39
			return fmt.Errorf("configured kv provider %s conflicts with %s storage in %q; restart with --kv-provider=%s to recover existing data, or use a separate empty --data-dir for a new store; do not remove the marker or stored data to bypass this check", option, existing, dir, existing)
40
		}
41
	}
42
	if option == "memory" {
43
		return nil
44
	}
45
	if marker != "" {
46
		// Retry a flush if an earlier startup published the marker but failed
47
		// to finish syncing its identity before returning an error.
48
		return syncKVProviderIdentity(dir)
49
	}
50
	return writeKVProviderMarker(dir, option)
51
}
52
 
53
func readKVProviderMarker(dir string) (string, error) {
54
	path := filepath.Join(dir, kvProviderMarker)
55
	info, err := os.Lstat(path)
56
	if os.IsNotExist(err) {
57
		return "", nil
58
	}
59
	if err != nil {
60
		return "", fmt.Errorf("inspecting storage provider marker: %w", err)
61
	}
62
	if !info.Mode().IsRegular() || info.Size() > 16 {
63
		return "", fmt.Errorf("invalid storage provider marker %q; preserve the data directory and recover its provider identity before restarting", path)
64
	}
65
	f, err := os.Open(path)
66
	if err != nil {
67
		return "", fmt.Errorf("reading storage provider marker: %w", err)
68
	}
69
	defer f.Close()
70
	data, err := io.ReadAll(io.LimitReader(f, 17))
71
	if err != nil {
72
		return "", fmt.Errorf("reading storage provider marker: %w", err)
73
	}
74
	switch string(data) {
75
	case "aof\n":
76
		return "aof", nil
77
	case "sqlite\n":
78
		return "sqlite", nil
79
	default:
80
		return "", fmt.Errorf("invalid storage provider marker %q; expected aof or sqlite followed by a newline; preserve the data directory and recover its provider identity before restarting", path)
81
	}
82
}
83
 
84
// inspectKVLayout only reads directory metadata. Opening a provider here could
85
// replay, repair, or create storage before we know the chosen provider is safe.
86
func inspectKVLayout(dir string) (string, error) {
87
	entries, err := os.ReadDir(dir)
88
	if os.IsNotExist(err) {
89
		return "", nil
90
	}
91
	if err != nil {
92
		return "", fmt.Errorf("reading data directory: %w", err)
93
	}
94
	var provider string
95
	for _, entry := range entries {
96
		// Inspect every existing vnode, including those beyond --virtual. Other
97
		// top-level entries (for example the SQLite runtime cache) aren't stores.
98
		if strings.Trim(entry.Name(), "0123456789") != "" {
99
			continue
100
		}
101
		vnodeDir := filepath.Join(dir, entry.Name())
102
		children, err := os.ReadDir(vnodeDir)
103
		if err != nil {
104
			return "", fmt.Errorf("reading vnode storage directory: %w", err)
105
		}
106
		for _, child := range children {
107
			var found string
108
			switch child.Name() {
109
			case "wal":
110
				found = "aof"
111
			case "sqlite3":
112
				found = "sqlite"
113
			default:
114
				return "", fmt.Errorf("unrecognized storage layout at %q; preserve the directory and recover or migrate the existing store before restarting", filepath.Join(vnodeDir, child.Name()))
115
			}
116
			info, err := os.Stat(filepath.Join(vnodeDir, child.Name()))
117
			if err != nil {
118
				return "", fmt.Errorf("inspecting %s storage directory: %w", found, err)
119
			}
120
			if !info.IsDir() {
121
				return "", fmt.Errorf("expected %s storage directory at %q", found, filepath.Join(vnodeDir, child.Name()))
122
			}
123
			if provider != "" && provider != found {
124
				return "", fmt.Errorf("mixed aof and sqlite storage in data directory %q; preserve the entire directory and resolve the stores through explicit recovery or migration before restarting", dir)
125
			}
126
			provider = found
127
		}
128
	}
129
	return provider, nil
130
}
131
 
132
func writeKVProviderMarker(dir, provider string) error {
133
	if err := prepareKVProviderDirectory(dir); err != nil {
134
		return fmt.Errorf("creating data directory: %w", err)
135
	}
136
	f, err := os.CreateTemp(dir, ".kv-provider-*")
137
	if err != nil {
138
		return fmt.Errorf("creating storage provider marker: %w", err)
139
	}
140
	defer os.Remove(f.Name())
141
	defer f.Close()
142
	if _, err := f.WriteString(provider + "\n"); err != nil {
143
		return fmt.Errorf("writing storage provider marker: %w", err)
144
	}
145
	if err := f.Sync(); err != nil {
146
		return fmt.Errorf("syncing storage provider marker: %w", err)
147
	}
148
	if err := f.Close(); err != nil {
149
		return fmt.Errorf("closing storage provider marker: %w", err)
150
	}
151
	if err := publishKVProviderMarker(f.Name(), dir); err != nil {
152
		return fmt.Errorf("publishing storage provider marker: %w", err)
153
	}
154
	return nil
155
}