fix(persistence): address PR review findings
- commit
- aead79d2e42b04fd3ca23041e32857d625c8bb13
- author
- brice · · 2026-07-03 20:51 UTC
- parent
- 6ff07bd
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
14 files · +237 -91
M internal/app/run.go +8 -3
| @@ -14,7 +14,6 @@ import ( | |||
| 14 | 14 | _ "net/http/pprof" | |
| 15 | 15 | "os" | |
| 16 | 16 | "os/signal" | |
| 17 | - | "path/filepath" | |
| 18 | 17 | "syscall" | |
| 19 | 18 | "text/tabwriter" | |
| 20 | 19 | "time" | |
| @@ -228,7 +227,7 @@ func runDaemon() error { | |||
| 228 | 227 | ||
| 229 | 228 | slog.Info("daemon starting", "socket", ln.Addr()) | |
| 230 | 229 | daemonOpts := []daemon.Option(nil) | |
| 231 | - | if store, err := kv.Open(filepath.Join(platform.StateDir(), "sessions.kv")); err != nil { | |
| 230 | + | if store, err := kv.Open(persist.StorePath(platform.StateDir())); err != nil { | |
| 232 | 231 | slog.Warn("opening session store failed; persistence disabled", "err", err) | |
| 233 | 232 | } else { | |
| 234 | 233 | daemonOpts = append(daemonOpts, daemon.WithStore(store)) | |
| @@ -560,7 +559,13 @@ func runKill(_ context.Context, name string, all, daemon bool) error { | |||
| 560 | 559 | transport, err := realDial(ipc.SocketDir()) | |
| 561 | 560 | if err != nil { | |
| 562 | 561 | if name != "" && !all && !daemon { | |
| 563 | - | store, openErr := kv.Open(filepath.Join(platform.StateDir(), "sessions.kv")) | |
| 562 | + | storePath := persist.StorePath(platform.StateDir()) | |
| 563 | + | if _, statErr := os.Stat(storePath); errors.Is(statErr, os.ErrNotExist) { | |
| 564 | + | return fmt.Errorf("vev: no such session: %s", name) | |
| 565 | + | } else if statErr != nil { | |
| 566 | + | return fmt.Errorf("vev: reading stored sessions: %w", statErr) | |
| 567 | + | } | |
| 568 | + | store, openErr := kv.Open(storePath) | |
| 564 | 569 | if openErr != nil { | |
| 565 | 570 | return fmt.Errorf("vev: opening stored sessions: %w", openErr) | |
| 566 | 571 | } | |
| @@ -14,7 +14,6 @@ import ( | |||||
| 14 | _ "net/http/pprof" | 14 | _ "net/http/pprof" | ||
| 15 | "os" | 15 | "os" | ||
| 16 | "os/signal" | 16 | "os/signal" | ||
| 17 | - | "path/filepath" | |||
| 18 | "syscall" | 17 | "syscall" | ||
| 19 | "text/tabwriter" | 18 | "text/tabwriter" | ||
| 20 | "time" | 19 | "time" | ||
| @@ -228,7 +227,7 @@ func runDaemon() error { | |||||
| 228 | 227 | ||||
| 229 | slog.Info("daemon starting", "socket", ln.Addr()) | 228 | slog.Info("daemon starting", "socket", ln.Addr()) | ||
| 230 | daemonOpts := []daemon.Option(nil) | 229 | daemonOpts := []daemon.Option(nil) | ||
| 231 | - | if store, err := kv.Open(filepath.Join(platform.StateDir(), "sessions.kv")); err != nil { | 230 | + | if store, err := kv.Open(persist.StorePath(platform.StateDir())); err != nil { |
| 232 | slog.Warn("opening session store failed; persistence disabled", "err", err) | 231 | slog.Warn("opening session store failed; persistence disabled", "err", err) | ||
| 233 | } else { | 232 | } else { | ||
| 234 | daemonOpts = append(daemonOpts, daemon.WithStore(store)) | 233 | daemonOpts = append(daemonOpts, daemon.WithStore(store)) | ||
| @@ -560,7 +559,13 @@ func runKill(_ context.Context, name string, all, daemon bool) error { | |||||
| 560 | transport, err := realDial(ipc.SocketDir()) | 559 | transport, err := realDial(ipc.SocketDir()) | ||
| 561 | if err != nil { | 560 | if err != nil { | ||
| 562 | if name != "" && !all && !daemon { | 561 | if name != "" && !all && !daemon { | ||
| 563 | - | store, openErr := kv.Open(filepath.Join(platform.StateDir(), "sessions.kv")) | 562 | + | storePath := persist.StorePath(platform.StateDir()) |
| 563 | + | if _, statErr := os.Stat(storePath); errors.Is(statErr, os.ErrNotExist) { | |||
| 564 | + | return fmt.Errorf("vev: no such session: %s", name) | |||
| 565 | + | } else if statErr != nil { | |||
| 566 | + | return fmt.Errorf("vev: reading stored sessions: %w", statErr) | |||
| 567 | + | } | |||
| 568 | + | store, openErr := kv.Open(storePath) | |||
| 564 | if openErr != nil { | 569 | if openErr != nil { | ||
| 565 | return fmt.Errorf("vev: opening stored sessions: %w", openErr) | 570 | return fmt.Errorf("vev: opening stored sessions: %w", openErr) | ||
| 566 | } | 571 | } | ||
M internal/app/run_test.go +14 -0
| @@ -144,6 +144,20 @@ func TestRunListReadsStoppedSessionsWithoutDaemon(t *testing.T) { | |||
| 144 | 144 | } | |
| 145 | 145 | } | |
| 146 | 146 | ||
| 147 | + | func TestRunKillMissingStoppedSessionDoesNotCreateStore(t *testing.T) { | |
| 148 | + | stateRoot, runtimeRoot := t.TempDir(), t.TempDir() | |
| 149 | + | t.Setenv("XDG_STATE_HOME", stateRoot) | |
| 150 | + | t.Setenv("XDG_RUNTIME_DIR", runtimeRoot) | |
| 151 | + | ||
| 152 | + | err := runKill(context.Background(), "missing", false, false) | |
| 153 | + | if err == nil || !strings.Contains(err.Error(), "no such session: missing") { | |
| 154 | + | t.Fatalf("runKill error = %v, want no such session", err) | |
| 155 | + | } | |
| 156 | + | if _, statErr := os.Stat(filepath.Join(stateRoot, "vev", "sessions.kv")); !errors.Is(statErr, os.ErrNotExist) { | |
| 157 | + | t.Fatalf("sessions.kv stat error = %v, want not exist", statErr) | |
| 158 | + | } | |
| 159 | + | } | |
| 160 | + | ||
| 147 | 161 | func TestRunKillDeletesStoppedSessionWithoutDaemon(t *testing.T) { | |
| 148 | 162 | stateRoot, runtimeRoot := t.TempDir(), t.TempDir() | |
| 149 | 163 | t.Setenv("XDG_STATE_HOME", stateRoot) | |
| @@ -144,6 +144,20 @@ func TestRunListReadsStoppedSessionsWithoutDaemon(t *testing.T) { | |||||
| 144 | } | 144 | } | ||
| 145 | } | 145 | } | ||
| 146 | 146 | ||||
| 147 | + | func TestRunKillMissingStoppedSessionDoesNotCreateStore(t *testing.T) { | |||
| 148 | + | stateRoot, runtimeRoot := t.TempDir(), t.TempDir() | |||
| 149 | + | t.Setenv("XDG_STATE_HOME", stateRoot) | |||
| 150 | + | t.Setenv("XDG_RUNTIME_DIR", runtimeRoot) | |||
| 151 | + | ||||
| 152 | + | err := runKill(context.Background(), "missing", false, false) | |||
| 153 | + | if err == nil || !strings.Contains(err.Error(), "no such session: missing") { | |||
| 154 | + | t.Fatalf("runKill error = %v, want no such session", err) | |||
| 155 | + | } | |||
| 156 | + | if _, statErr := os.Stat(filepath.Join(stateRoot, "vev", "sessions.kv")); !errors.Is(statErr, os.ErrNotExist) { | |||
| 157 | + | t.Fatalf("sessions.kv stat error = %v, want not exist", statErr) | |||
| 158 | + | } | |||
| 159 | + | } | |||
| 160 | + | ||||
| 147 | func TestRunKillDeletesStoppedSessionWithoutDaemon(t *testing.T) { | 161 | func TestRunKillDeletesStoppedSessionWithoutDaemon(t *testing.T) { | ||
| 148 | stateRoot, runtimeRoot := t.TempDir(), t.TempDir() | 162 | stateRoot, runtimeRoot := t.TempDir(), t.TempDir() | ||
| 149 | t.Setenv("XDG_STATE_HOME", stateRoot) | 163 | t.Setenv("XDG_STATE_HOME", stateRoot) | ||
M internal/persist/persist.go +13 -11
| @@ -1,6 +1,7 @@ | |||
| 1 | 1 | package persist | |
| 2 | 2 | ||
| 3 | 3 | import ( | |
| 4 | + | "log/slog" | |
| 4 | 5 | "path/filepath" | |
| 5 | 6 | "sort" | |
| 6 | 7 | ||
| @@ -10,6 +11,11 @@ import ( | |||
| 10 | 11 | ||
| 11 | 12 | const filename = "sessions.kv" | |
| 12 | 13 | ||
| 14 | + | // StorePath returns the canonical path for the persisted session metadata store. | |
| 15 | + | func StorePath(dir string) string { | |
| 16 | + | return filepath.Join(dir, filename) | |
| 17 | + | } | |
| 18 | + | ||
| 13 | 19 | // Persister stores named session metadata. A nil Store makes all mutating | |
| 14 | 20 | // operations no-ops, so callers can keep persistence optional. | |
| 15 | 21 | type Persister struct { | |
| @@ -18,7 +24,7 @@ type Persister struct { | |||
| 18 | 24 | ||
| 19 | 25 | // Open opens the session persister under dir. | |
| 20 | 26 | func Open(dir string) (*Persister, error) { | |
| 21 | - | store, err := kv.Open(filepath.Join(dir, filename)) | |
| 27 | + | store, err := kv.Open(StorePath(dir)) | |
| 22 | 28 | if err != nil { | |
| 23 | 29 | return nil, err | |
| 24 | 30 | } | |
| @@ -32,7 +38,7 @@ func New(store ports.Store) *Persister { | |||
| 32 | 38 | ||
| 33 | 39 | // LoadReadOnly replays the session store under dir without mutating it. | |
| 34 | 40 | func LoadReadOnly(dir string) ([]Record, error) { | |
| 35 | - | data, err := kv.Replay(filepath.Join(dir, filename)) | |
| 41 | + | data, err := kv.Replay(StorePath(dir)) | |
| 36 | 42 | if err != nil { | |
| 37 | 43 | return nil, err | |
| 38 | 44 | } | |
| @@ -64,16 +70,11 @@ func (p *Persister) Touch(name, cwd string, at int64) error { | |||
| 64 | 70 | } | |
| 65 | 71 | ||
| 66 | 72 | createdAt := at | |
| 67 | - | p.store.Range(func(k, v []byte) bool { | |
| 68 | - | if string(k) != name { | |
| 69 | - | return true | |
| 70 | - | } | |
| 71 | - | r, err := decodeRecordValue(name, v) | |
| 72 | - | if err == nil { | |
| 73 | + | if v, ok := p.store.Get([]byte(name)); ok { | |
| 74 | + | if r, err := decodeRecordValue(name, v); err == nil { | |
| 73 | 75 | createdAt = r.CreatedAt | |
| 74 | 76 | } | |
| 75 | - | return false | |
| 76 | - | }) | |
| 77 | + | } | |
| 77 | 78 | ||
| 78 | 79 | value, err := encodeRecordValue(Record{Name: name, Cwd: cwd, CreatedAt: createdAt, UpdatedAt: at}) | |
| 79 | 80 | if err != nil { | |
| @@ -122,7 +123,8 @@ func decodeAll(data map[string][]byte) ([]Record, error) { | |||
| 122 | 123 | for name, value := range data { | |
| 123 | 124 | r, err := decodeRecordValue(name, value) | |
| 124 | 125 | if err != nil { | |
| 125 | - | return nil, err | |
| 126 | + | slog.Warn("skipping malformed persisted session", "session", name, "err", err) | |
| 127 | + | continue | |
| 126 | 128 | } | |
| 127 | 129 | records = append(records, r) | |
| 128 | 130 | } | |
| @@ -1,6 +1,7 @@ | |||||
| 1 | package persist | 1 | package persist | ||
| 2 | 2 | ||||
| 3 | import ( | 3 | import ( | ||
| 4 | + | "log/slog" | |||
| 4 | "path/filepath" | 5 | "path/filepath" | ||
| 5 | "sort" | 6 | "sort" | ||
| 6 | 7 | ||||
| @@ -10,6 +11,11 @@ import ( | |||||
| 10 | 11 | ||||
| 11 | const filename = "sessions.kv" | 12 | const filename = "sessions.kv" | ||
| 12 | 13 | ||||
| 14 | + | // StorePath returns the canonical path for the persisted session metadata store. | |||
| 15 | + | func StorePath(dir string) string { | |||
| 16 | + | return filepath.Join(dir, filename) | |||
| 17 | + | } | |||
| 18 | + | ||||
| 13 | // Persister stores named session metadata. A nil Store makes all mutating | 19 | // Persister stores named session metadata. A nil Store makes all mutating | ||
| 14 | // operations no-ops, so callers can keep persistence optional. | 20 | // operations no-ops, so callers can keep persistence optional. | ||
| 15 | type Persister struct { | 21 | type Persister struct { | ||
| @@ -18,7 +24,7 @@ type Persister struct { | |||||
| 18 | 24 | ||||
| 19 | // Open opens the session persister under dir. | 25 | // Open opens the session persister under dir. | ||
| 20 | func Open(dir string) (*Persister, error) { | 26 | func Open(dir string) (*Persister, error) { | ||
| 21 | - | store, err := kv.Open(filepath.Join(dir, filename)) | 27 | + | store, err := kv.Open(StorePath(dir)) |
| 22 | if err != nil { | 28 | if err != nil { | ||
| 23 | return nil, err | 29 | return nil, err | ||
| 24 | } | 30 | } | ||
| @@ -32,7 +38,7 @@ func New(store ports.Store) *Persister { | |||||
| 32 | 38 | ||||
| 33 | // LoadReadOnly replays the session store under dir without mutating it. | 39 | // LoadReadOnly replays the session store under dir without mutating it. | ||
| 34 | func LoadReadOnly(dir string) ([]Record, error) { | 40 | func LoadReadOnly(dir string) ([]Record, error) { | ||
| 35 | - | data, err := kv.Replay(filepath.Join(dir, filename)) | 41 | + | data, err := kv.Replay(StorePath(dir)) |
| 36 | if err != nil { | 42 | if err != nil { | ||
| 37 | return nil, err | 43 | return nil, err | ||
| 38 | } | 44 | } | ||
| @@ -64,16 +70,11 @@ func (p *Persister) Touch(name, cwd string, at int64) error { | |||||
| 64 | } | 70 | } | ||
| 65 | 71 | ||||
| 66 | createdAt := at | 72 | createdAt := at | ||
| 67 | - | p.store.Range(func(k, v []byte) bool { | 73 | + | if v, ok := p.store.Get([]byte(name)); ok { |
| 68 | - | if string(k) != name { | 74 | + | if r, err := decodeRecordValue(name, v); err == nil { |
| 69 | - | return true | |||
| 70 | - | } | |||
| 71 | - | r, err := decodeRecordValue(name, v) | |||
| 72 | - | if err == nil { | |||
| 73 | createdAt = r.CreatedAt | 75 | createdAt = r.CreatedAt | ||
| 74 | } | 76 | } | ||
| 75 | - | return false | 77 | + | } |
| 76 | - | }) | |||
| 77 | 78 | ||||
| 78 | value, err := encodeRecordValue(Record{Name: name, Cwd: cwd, CreatedAt: createdAt, UpdatedAt: at}) | 79 | value, err := encodeRecordValue(Record{Name: name, Cwd: cwd, CreatedAt: createdAt, UpdatedAt: at}) | ||
| 79 | if err != nil { | 80 | if err != nil { | ||
| @@ -122,7 +123,8 @@ func decodeAll(data map[string][]byte) ([]Record, error) { | |||||
| 122 | for name, value := range data { | 123 | for name, value := range data { | ||
| 123 | r, err := decodeRecordValue(name, value) | 124 | r, err := decodeRecordValue(name, value) | ||
| 124 | if err != nil { | 125 | if err != nil { | ||
| 125 | - | return nil, err | 126 | + | slog.Warn("skipping malformed persisted session", "session", name, "err", err) |
| 127 | + | continue | |||
| 126 | } | 128 | } | ||
| 127 | records = append(records, r) | 129 | records = append(records, r) | ||
| 128 | } | 130 | } | ||
M internal/persist/persist_test.go +63 -46
| @@ -5,7 +5,10 @@ import ( | |||
| 5 | 5 | "path/filepath" | |
| 6 | 6 | "testing" | |
| 7 | 7 | ||
| 8 | + | "github.com/stretchr/testify/mock" | |
| 8 | 9 | "github.com/stretchr/testify/require" | |
| 10 | + | ||
| 11 | + | portsmocks "github.com/bnema/vev/internal/ports/mocks" | |
| 9 | 12 | ) | |
| 10 | 13 | ||
| 11 | 14 | func TestRecordCodecRoundTrip(t *testing.T) { | |
| @@ -82,6 +85,18 @@ func TestPersisterSaveTouchDeleteLoadAll(t *testing.T) { | |||
| 82 | 85 | require.NoError(t, p.Close()) | |
| 83 | 86 | } | |
| 84 | 87 | ||
| 88 | + | func TestDecodeAllSkipsMalformedRecords(t *testing.T) { | |
| 89 | + | good, err := encodeRecordValue(Record{Name: "good", Cwd: "/ok", CreatedAt: 1, UpdatedAt: 2}) | |
| 90 | + | require.NoError(t, err) | |
| 91 | + | ||
| 92 | + | records, err := decodeAll(map[string][]byte{ | |
| 93 | + | "bad": []byte{0, 0, 0, 5, 'x'}, | |
| 94 | + | "good": good, | |
| 95 | + | }) | |
| 96 | + | require.NoError(t, err) | |
| 97 | + | require.Equal(t, []Record{{Name: "good", Cwd: "/ok", CreatedAt: 1, UpdatedAt: 2}}, records) | |
| 98 | + | } | |
| 99 | + | ||
| 85 | 100 | func TestLoadReadOnlyUsesReplay(t *testing.T) { | |
| 86 | 101 | dir := t.TempDir() | |
| 87 | 102 | p, err := Open(dir) | |
| @@ -101,33 +116,33 @@ func TestLoadReadOnlyUsesReplay(t *testing.T) { | |||
| 101 | 116 | } | |
| 102 | 117 | ||
| 103 | 118 | func TestSyncBehavior(t *testing.T) { | |
| 104 | - | store := newFakeStore() | |
| 119 | + | store, state := newMockStore(t) | |
| 105 | 120 | p := New(store) | |
| 106 | 121 | ||
| 107 | 122 | require.NoError(t, p.Save(Record{Name: "work", Cwd: "/work", CreatedAt: 1, UpdatedAt: 1})) | |
| 108 | - | require.Equal(t, 1, store.sets) | |
| 109 | - | require.Equal(t, 1, store.syncs) | |
| 123 | + | require.Equal(t, 1, state.sets) | |
| 124 | + | require.Equal(t, 1, state.syncs) | |
| 110 | 125 | ||
| 111 | 126 | require.NoError(t, p.Touch("work", "/work/next", 2)) | |
| 112 | - | require.Equal(t, 2, store.sets) | |
| 113 | - | require.Equal(t, 1, store.syncs) | |
| 127 | + | require.Equal(t, 2, state.sets) | |
| 128 | + | require.Equal(t, 1, state.syncs) | |
| 114 | 129 | ||
| 115 | 130 | require.NoError(t, p.Delete("work")) | |
| 116 | - | require.Equal(t, 1, store.deletes) | |
| 117 | - | require.Equal(t, 2, store.syncs) | |
| 131 | + | require.Equal(t, 1, state.deletes) | |
| 132 | + | require.Equal(t, 2, state.syncs) | |
| 118 | 133 | } | |
| 119 | 134 | ||
| 120 | 135 | func TestSavePropagatesSetErrorWithoutSync(t *testing.T) { | |
| 121 | - | store := newFakeStore() | |
| 122 | - | store.errSet = errors.New("set failed") | |
| 136 | + | store, state := newMockStore(t) | |
| 137 | + | state.errSet = errors.New("set failed") | |
| 123 | 138 | p := New(store) | |
| 124 | 139 | ||
| 125 | 140 | err := p.Save(Record{Name: "work", Cwd: "/work", CreatedAt: 1, UpdatedAt: 1}) | |
| 126 | - | require.ErrorIs(t, err, store.errSet) | |
| 127 | - | require.Zero(t, store.syncs) | |
| 141 | + | require.ErrorIs(t, err, state.errSet) | |
| 142 | + | require.Zero(t, state.syncs) | |
| 128 | 143 | } | |
| 129 | 144 | ||
| 130 | - | type fakeStore struct { | |
| 145 | + | type mockStoreState struct { | |
| 131 | 146 | data map[string][]byte | |
| 132 | 147 | sets int | |
| 133 | 148 | deletes int | |
| @@ -136,39 +151,41 @@ type fakeStore struct { | |||
| 136 | 151 | errSet error | |
| 137 | 152 | } | |
| 138 | 153 | ||
| 139 | - | func newFakeStore() *fakeStore { | |
| 140 | - | return &fakeStore{data: make(map[string][]byte)} | |
| 141 | - | } | |
| 142 | - | ||
| 143 | - | func (s *fakeStore) Set(key, val []byte) error { | |
| 144 | - | s.sets++ | |
| 145 | - | if s.errSet != nil { | |
| 146 | - | return s.errSet | |
| 147 | - | } | |
| 148 | - | s.data[string(key)] = append([]byte(nil), val...) | |
| 149 | - | return nil | |
| 150 | - | } | |
| 151 | - | ||
| 152 | - | func (s *fakeStore) Delete(key []byte) error { | |
| 153 | - | s.deletes++ | |
| 154 | - | delete(s.data, string(key)) | |
| 155 | - | return nil | |
| 156 | - | } | |
| 157 | - | ||
| 158 | - | func (s *fakeStore) Range(fn func(k, v []byte) bool) { | |
| 159 | - | for k, v := range s.data { | |
| 160 | - | if !fn([]byte(k), append([]byte(nil), v...)) { | |
| 161 | - | return | |
| 154 | + | func newMockStore(t *testing.T) (*portsmocks.MockStore, *mockStoreState) { | |
| 155 | + | t.Helper() | |
| 156 | + | state := &mockStoreState{data: make(map[string][]byte)} | |
| 157 | + | store := portsmocks.NewMockStore(t) | |
| 158 | + | store.EXPECT().Get(mock.Anything).RunAndReturn(func(key []byte) ([]byte, bool) { | |
| 159 | + | v, ok := state.data[string(key)] | |
| 160 | + | return append([]byte(nil), v...), ok | |
| 161 | + | }).Maybe() | |
| 162 | + | store.EXPECT().Set(mock.Anything, mock.Anything).RunAndReturn(func(key, val []byte) error { | |
| 163 | + | state.sets++ | |
| 164 | + | if state.errSet != nil { | |
| 165 | + | return state.errSet | |
| 162 | 166 | } | |
| 163 | - | } | |
| 164 | - | } | |
| 165 | - | ||
| 166 | - | func (s *fakeStore) Sync() error { | |
| 167 | - | s.syncs++ | |
| 168 | - | return nil | |
| 169 | - | } | |
| 170 | - | ||
| 171 | - | func (s *fakeStore) Close() error { | |
| 172 | - | s.closed = true | |
| 173 | - | return nil | |
| 167 | + | state.data[string(key)] = append([]byte(nil), val...) | |
| 168 | + | return nil | |
| 169 | + | }).Maybe() | |
| 170 | + | store.EXPECT().Delete(mock.Anything).RunAndReturn(func(key []byte) error { | |
| 171 | + | state.deletes++ | |
| 172 | + | delete(state.data, string(key)) | |
| 173 | + | return nil | |
| 174 | + | }).Maybe() | |
| 175 | + | store.EXPECT().Range(mock.Anything).Run(func(fn func(k, v []byte) bool) { | |
| 176 | + | for k, v := range state.data { | |
| 177 | + | if !fn([]byte(k), append([]byte(nil), v...)) { | |
| 178 | + | return | |
| 179 | + | } | |
| 180 | + | } | |
| 181 | + | }).Maybe() | |
| 182 | + | store.EXPECT().Sync().RunAndReturn(func() error { | |
| 183 | + | state.syncs++ | |
| 184 | + | return nil | |
| 185 | + | }).Maybe() | |
| 186 | + | store.EXPECT().Close().RunAndReturn(func() error { | |
| 187 | + | state.closed = true | |
| 188 | + | return nil | |
| 189 | + | }).Maybe() | |
| 190 | + | return store, state | |
| 174 | 191 | } | |
| @@ -5,7 +5,10 @@ import ( | |||||
| 5 | "path/filepath" | 5 | "path/filepath" | ||
| 6 | "testing" | 6 | "testing" | ||
| 7 | 7 | ||||
| 8 | + | "github.com/stretchr/testify/mock" | |||
| 8 | "github.com/stretchr/testify/require" | 9 | "github.com/stretchr/testify/require" | ||
| 10 | + | ||||
| 11 | + | portsmocks "github.com/bnema/vev/internal/ports/mocks" | |||
| 9 | ) | 12 | ) | ||
| 10 | 13 | ||||
| 11 | func TestRecordCodecRoundTrip(t *testing.T) { | 14 | func TestRecordCodecRoundTrip(t *testing.T) { | ||
| @@ -82,6 +85,18 @@ func TestPersisterSaveTouchDeleteLoadAll(t *testing.T) { | |||||
| 82 | require.NoError(t, p.Close()) | 85 | require.NoError(t, p.Close()) | ||
| 83 | } | 86 | } | ||
| 84 | 87 | ||||
| 88 | + | func TestDecodeAllSkipsMalformedRecords(t *testing.T) { | |||
| 89 | + | good, err := encodeRecordValue(Record{Name: "good", Cwd: "/ok", CreatedAt: 1, UpdatedAt: 2}) | |||
| 90 | + | require.NoError(t, err) | |||
| 91 | + | ||||
| 92 | + | records, err := decodeAll(map[string][]byte{ | |||
| 93 | + | "bad": []byte{0, 0, 0, 5, 'x'}, | |||
| 94 | + | "good": good, | |||
| 95 | + | }) | |||
| 96 | + | require.NoError(t, err) | |||
| 97 | + | require.Equal(t, []Record{{Name: "good", Cwd: "/ok", CreatedAt: 1, UpdatedAt: 2}}, records) | |||
| 98 | + | } | |||
| 99 | + | ||||
| 85 | func TestLoadReadOnlyUsesReplay(t *testing.T) { | 100 | func TestLoadReadOnlyUsesReplay(t *testing.T) { | ||
| 86 | dir := t.TempDir() | 101 | dir := t.TempDir() | ||
| 87 | p, err := Open(dir) | 102 | p, err := Open(dir) | ||
| @@ -101,33 +116,33 @@ func TestLoadReadOnlyUsesReplay(t *testing.T) { | |||||
| 101 | } | 116 | } | ||
| 102 | 117 | ||||
| 103 | func TestSyncBehavior(t *testing.T) { | 118 | func TestSyncBehavior(t *testing.T) { | ||
| 104 | - | store := newFakeStore() | 119 | + | store, state := newMockStore(t) |
| 105 | p := New(store) | 120 | p := New(store) | ||
| 106 | 121 | ||||
| 107 | require.NoError(t, p.Save(Record{Name: "work", Cwd: "/work", CreatedAt: 1, UpdatedAt: 1})) | 122 | require.NoError(t, p.Save(Record{Name: "work", Cwd: "/work", CreatedAt: 1, UpdatedAt: 1})) | ||
| 108 | - | require.Equal(t, 1, store.sets) | 123 | + | require.Equal(t, 1, state.sets) |
| 109 | - | require.Equal(t, 1, store.syncs) | 124 | + | require.Equal(t, 1, state.syncs) |
| 110 | 125 | ||||
| 111 | require.NoError(t, p.Touch("work", "/work/next", 2)) | 126 | require.NoError(t, p.Touch("work", "/work/next", 2)) | ||
| 112 | - | require.Equal(t, 2, store.sets) | 127 | + | require.Equal(t, 2, state.sets) |
| 113 | - | require.Equal(t, 1, store.syncs) | 128 | + | require.Equal(t, 1, state.syncs) |
| 114 | 129 | ||||
| 115 | require.NoError(t, p.Delete("work")) | 130 | require.NoError(t, p.Delete("work")) | ||
| 116 | - | require.Equal(t, 1, store.deletes) | 131 | + | require.Equal(t, 1, state.deletes) |
| 117 | - | require.Equal(t, 2, store.syncs) | 132 | + | require.Equal(t, 2, state.syncs) |
| 118 | } | 133 | } | ||
| 119 | 134 | ||||
| 120 | func TestSavePropagatesSetErrorWithoutSync(t *testing.T) { | 135 | func TestSavePropagatesSetErrorWithoutSync(t *testing.T) { | ||
| 121 | - | store := newFakeStore() | 136 | + | store, state := newMockStore(t) |
| 122 | - | store.errSet = errors.New("set failed") | 137 | + | state.errSet = errors.New("set failed") |
| 123 | p := New(store) | 138 | p := New(store) | ||
| 124 | 139 | ||||
| 125 | err := p.Save(Record{Name: "work", Cwd: "/work", CreatedAt: 1, UpdatedAt: 1}) | 140 | err := p.Save(Record{Name: "work", Cwd: "/work", CreatedAt: 1, UpdatedAt: 1}) | ||
| 126 | - | require.ErrorIs(t, err, store.errSet) | 141 | + | require.ErrorIs(t, err, state.errSet) |
| 127 | - | require.Zero(t, store.syncs) | 142 | + | require.Zero(t, state.syncs) |
| 128 | } | 143 | } | ||
| 129 | 144 | ||||
| 130 | - | type fakeStore struct { | 145 | + | type mockStoreState struct { |
| 131 | data map[string][]byte | 146 | data map[string][]byte | ||
| 132 | sets int | 147 | sets int | ||
| 133 | deletes int | 148 | deletes int | ||
| @@ -136,39 +151,41 @@ type fakeStore struct { | |||||
| 136 | errSet error | 151 | errSet error | ||
| 137 | } | 152 | } | ||
| 138 | 153 | ||||
| 139 | - | func newFakeStore() *fakeStore { | 154 | + | func newMockStore(t *testing.T) (*portsmocks.MockStore, *mockStoreState) { |
| 140 | - | return &fakeStore{data: make(map[string][]byte)} | 155 | + | t.Helper() |
| 141 | - | } | 156 | + | state := &mockStoreState{data: make(map[string][]byte)} |
| 142 | - | 157 | + | store := portsmocks.NewMockStore(t) | |
| 143 | - | func (s *fakeStore) Set(key, val []byte) error { | 158 | + | store.EXPECT().Get(mock.Anything).RunAndReturn(func(key []byte) ([]byte, bool) { |
| 144 | - | s.sets++ | 159 | + | v, ok := state.data[string(key)] |
| 145 | - | if s.errSet != nil { | 160 | + | return append([]byte(nil), v...), ok |
| 146 | - | return s.errSet | 161 | + | }).Maybe() |
| 147 | - | } | 162 | + | store.EXPECT().Set(mock.Anything, mock.Anything).RunAndReturn(func(key, val []byte) error { |
| 148 | - | s.data[string(key)] = append([]byte(nil), val...) | 163 | + | state.sets++ |
| 149 | - | return nil | 164 | + | if state.errSet != nil { |
| 150 | - | } | 165 | + | return state.errSet |
| 151 | - | ||||
| 152 | - | func (s *fakeStore) Delete(key []byte) error { | |||
| 153 | - | s.deletes++ | |||
| 154 | - | delete(s.data, string(key)) | |||
| 155 | - | return nil | |||
| 156 | - | } | |||
| 157 | - | ||||
| 158 | - | func (s *fakeStore) Range(fn func(k, v []byte) bool) { | |||
| 159 | - | for k, v := range s.data { | |||
| 160 | - | if !fn([]byte(k), append([]byte(nil), v...)) { | |||
| 161 | - | return | |||
| 162 | } | 166 | } | ||
| 163 | - | } | 167 | + | state.data[string(key)] = append([]byte(nil), val...) |
| 164 | - | } | 168 | + | return nil |
| 165 | - | 169 | + | }).Maybe() | |
| 166 | - | func (s *fakeStore) Sync() error { | 170 | + | store.EXPECT().Delete(mock.Anything).RunAndReturn(func(key []byte) error { |
| 167 | - | s.syncs++ | 171 | + | state.deletes++ |
| 168 | - | return nil | 172 | + | delete(state.data, string(key)) |
| 169 | - | } | 173 | + | return nil |
| 170 | - | 174 | + | }).Maybe() | |
| 171 | - | func (s *fakeStore) Close() error { | 175 | + | store.EXPECT().Range(mock.Anything).Run(func(fn func(k, v []byte) bool) { |
| 172 | - | s.closed = true | 176 | + | for k, v := range state.data { |
| 173 | - | return nil | 177 | + | if !fn([]byte(k), append([]byte(nil), v...)) { |
| 178 | + | return | |||
| 179 | + | } | |||
| 180 | + | } | |||
| 181 | + | }).Maybe() | |||
| 182 | + | store.EXPECT().Sync().RunAndReturn(func() error { | |||
| 183 | + | state.syncs++ | |||
| 184 | + | return nil | |||
| 185 | + | }).Maybe() | |||
| 186 | + | store.EXPECT().Close().RunAndReturn(func() error { | |||
| 187 | + | state.closed = true | |||
| 188 | + | return nil | |||
| 189 | + | }).Maybe() | |||
| 190 | + | return store, state | |||
| 174 | } | 191 | } | ||
M internal/persist/record.go +2 -1
| @@ -4,6 +4,7 @@ import ( | |||
| 4 | 4 | "encoding/binary" | |
| 5 | 5 | "errors" | |
| 6 | 6 | "fmt" | |
| 7 | + | "math" | |
| 7 | 8 | ) | |
| 8 | 9 | ||
| 9 | 10 | // Record is the persisted metadata for a named session. | |
| @@ -23,7 +24,7 @@ func encodeRecordValue(r Record) ([]byte, error) { | |||
| 23 | 24 | if r.Name == "" { | |
| 24 | 25 | return nil, errEmptyName | |
| 25 | 26 | } | |
| 26 | - | if len(r.Cwd) > int(^uint32(0)) { | |
| 27 | + | if uint64(len(r.Cwd)) > math.MaxUint32 { | |
| 27 | 28 | return nil, fmt.Errorf("persist: cwd too large") | |
| 28 | 29 | } | |
| 29 | 30 | ||
| @@ -4,6 +4,7 @@ import ( | |||||
| 4 | "encoding/binary" | 4 | "encoding/binary" | ||
| 5 | "errors" | 5 | "errors" | ||
| 6 | "fmt" | 6 | "fmt" | ||
| 7 | + | "math" | |||
| 7 | ) | 8 | ) | ||
| 8 | 9 | ||||
| 9 | // Record is the persisted metadata for a named session. | 10 | // Record is the persisted metadata for a named session. | ||
| @@ -23,7 +24,7 @@ func encodeRecordValue(r Record) ([]byte, error) { | |||||
| 23 | if r.Name == "" { | 24 | if r.Name == "" { | ||
| 24 | return nil, errEmptyName | 25 | return nil, errEmptyName | ||
| 25 | } | 26 | } | ||
| 26 | - | if len(r.Cwd) > int(^uint32(0)) { | 27 | + | if uint64(len(r.Cwd)) > math.MaxUint32 { |
| 27 | return nil, fmt.Errorf("persist: cwd too large") | 28 | return nil, fmt.Errorf("persist: cwd too large") | ||
| 28 | } | 29 | } | ||
| 29 | 30 | ||||
M internal/ports/mocks/mocks.go +62 -0
| @@ -1196,6 +1196,68 @@ func (_c *MockStore_Delete_Call) RunAndReturn(run func(key []byte) error) *MockS | |||
| 1196 | 1196 | return _c | |
| 1197 | 1197 | } | |
| 1198 | 1198 | ||
| 1199 | + | // Get provides a mock function for the type MockStore | |
| 1200 | + | func (_mock *MockStore) Get(key []byte) ([]byte, bool) { | |
| 1201 | + | ret := _mock.Called(key) | |
| 1202 | + | ||
| 1203 | + | if len(ret) == 0 { | |
| 1204 | + | panic("no return value specified for Get") | |
| 1205 | + | } | |
| 1206 | + | ||
| 1207 | + | var r0 []byte | |
| 1208 | + | var r1 bool | |
| 1209 | + | if returnFunc, ok := ret.Get(0).(func([]byte) ([]byte, bool)); ok { | |
| 1210 | + | return returnFunc(key) | |
| 1211 | + | } | |
| 1212 | + | if returnFunc, ok := ret.Get(0).(func([]byte) []byte); ok { | |
| 1213 | + | r0 = returnFunc(key) | |
| 1214 | + | } else { | |
| 1215 | + | if ret.Get(0) != nil { | |
| 1216 | + | r0 = ret.Get(0).([]byte) | |
| 1217 | + | } | |
| 1218 | + | } | |
| 1219 | + | if returnFunc, ok := ret.Get(1).(func([]byte) bool); ok { | |
| 1220 | + | r1 = returnFunc(key) | |
| 1221 | + | } else { | |
| 1222 | + | r1 = ret.Get(1).(bool) | |
| 1223 | + | } | |
| 1224 | + | return r0, r1 | |
| 1225 | + | } | |
| 1226 | + | ||
| 1227 | + | // MockStore_Get_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Get' | |
| 1228 | + | type MockStore_Get_Call struct { | |
| 1229 | + | *mock.Call | |
| 1230 | + | } | |
| 1231 | + | ||
| 1232 | + | // Get is a helper method to define mock.On call | |
| 1233 | + | // - key []byte | |
| 1234 | + | func (_e *MockStore_Expecter) Get(key any) *MockStore_Get_Call { | |
| 1235 | + | return &MockStore_Get_Call{Call: _e.mock.On("Get", key)} | |
| 1236 | + | } | |
| 1237 | + | ||
| 1238 | + | func (_c *MockStore_Get_Call) Run(run func(key []byte)) *MockStore_Get_Call { | |
| 1239 | + | _c.Call.Run(func(args mock.Arguments) { | |
| 1240 | + | var arg0 []byte | |
| 1241 | + | if args[0] != nil { | |
| 1242 | + | arg0 = args[0].([]byte) | |
| 1243 | + | } | |
| 1244 | + | run( | |
| 1245 | + | arg0, | |
| 1246 | + | ) | |
| 1247 | + | }) | |
| 1248 | + | return _c | |
| 1249 | + | } | |
| 1250 | + | ||
| 1251 | + | func (_c *MockStore_Get_Call) Return(bytes []byte, b bool) *MockStore_Get_Call { | |
| 1252 | + | _c.Call.Return(bytes, b) | |
| 1253 | + | return _c | |
| 1254 | + | } | |
| 1255 | + | ||
| 1256 | + | func (_c *MockStore_Get_Call) RunAndReturn(run func(key []byte) ([]byte, bool)) *MockStore_Get_Call { | |
| 1257 | + | _c.Call.Return(run) | |
| 1258 | + | return _c | |
| 1259 | + | } | |
| 1260 | + | ||
| 1199 | 1261 | // Range provides a mock function for the type MockStore | |
| 1200 | 1262 | func (_mock *MockStore) Range(fn func(k []byte, v []byte) bool) { | |
| 1201 | 1263 | _mock.Called(fn) | |
| @@ -1196,6 +1196,68 @@ func (_c *MockStore_Delete_Call) RunAndReturn(run func(key []byte) error) *MockS | |||||
| 1196 | return _c | 1196 | return _c | ||
| 1197 | } | 1197 | } | ||
| 1198 | 1198 | ||||
| 1199 | + | // Get provides a mock function for the type MockStore | |||
| 1200 | + | func (_mock *MockStore) Get(key []byte) ([]byte, bool) { | |||
| 1201 | + | ret := _mock.Called(key) | |||
| 1202 | + | ||||
| 1203 | + | if len(ret) == 0 { | |||
| 1204 | + | panic("no return value specified for Get") | |||
| 1205 | + | } | |||
| 1206 | + | ||||
| 1207 | + | var r0 []byte | |||
| 1208 | + | var r1 bool | |||
| 1209 | + | if returnFunc, ok := ret.Get(0).(func([]byte) ([]byte, bool)); ok { | |||
| 1210 | + | return returnFunc(key) | |||
| 1211 | + | } | |||
| 1212 | + | if returnFunc, ok := ret.Get(0).(func([]byte) []byte); ok { | |||
| 1213 | + | r0 = returnFunc(key) | |||
| 1214 | + | } else { | |||
| 1215 | + | if ret.Get(0) != nil { | |||
| 1216 | + | r0 = ret.Get(0).([]byte) | |||
| 1217 | + | } | |||
| 1218 | + | } | |||
| 1219 | + | if returnFunc, ok := ret.Get(1).(func([]byte) bool); ok { | |||
| 1220 | + | r1 = returnFunc(key) | |||
| 1221 | + | } else { | |||
| 1222 | + | r1 = ret.Get(1).(bool) | |||
| 1223 | + | } | |||
| 1224 | + | return r0, r1 | |||
| 1225 | + | } | |||
| 1226 | + | ||||
| 1227 | + | // MockStore_Get_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Get' | |||
| 1228 | + | type MockStore_Get_Call struct { | |||
| 1229 | + | *mock.Call | |||
| 1230 | + | } | |||
| 1231 | + | ||||
| 1232 | + | // Get is a helper method to define mock.On call | |||
| 1233 | + | // - key []byte | |||
| 1234 | + | func (_e *MockStore_Expecter) Get(key any) *MockStore_Get_Call { | |||
| 1235 | + | return &MockStore_Get_Call{Call: _e.mock.On("Get", key)} | |||
| 1236 | + | } | |||
| 1237 | + | ||||
| 1238 | + | func (_c *MockStore_Get_Call) Run(run func(key []byte)) *MockStore_Get_Call { | |||
| 1239 | + | _c.Call.Run(func(args mock.Arguments) { | |||
| 1240 | + | var arg0 []byte | |||
| 1241 | + | if args[0] != nil { | |||
| 1242 | + | arg0 = args[0].([]byte) | |||
| 1243 | + | } | |||
| 1244 | + | run( | |||
| 1245 | + | arg0, | |||
| 1246 | + | ) | |||
| 1247 | + | }) | |||
| 1248 | + | return _c | |||
| 1249 | + | } | |||
| 1250 | + | ||||
| 1251 | + | func (_c *MockStore_Get_Call) Return(bytes []byte, b bool) *MockStore_Get_Call { | |||
| 1252 | + | _c.Call.Return(bytes, b) | |||
| 1253 | + | return _c | |||
| 1254 | + | } | |||
| 1255 | + | ||||
| 1256 | + | func (_c *MockStore_Get_Call) RunAndReturn(run func(key []byte) ([]byte, bool)) *MockStore_Get_Call { | |||
| 1257 | + | _c.Call.Return(run) | |||
| 1258 | + | return _c | |||
| 1259 | + | } | |||
| 1260 | + | ||||
| 1199 | // Range provides a mock function for the type MockStore | 1261 | // Range provides a mock function for the type MockStore | ||
| 1200 | func (_mock *MockStore) Range(fn func(k []byte, v []byte) bool) { | 1262 | func (_mock *MockStore) Range(fn func(k []byte, v []byte) bool) { | ||
| 1201 | _mock.Called(fn) | 1263 | _mock.Called(fn) | ||
M internal/ports/ports.go +1 -0
| @@ -47,6 +47,7 @@ type Listener interface { | |||
| 47 | 47 | // | |
| 48 | 48 | // Implementations may buffer writes; Sync is the durability barrier. | |
| 49 | 49 | type Store interface { | |
| 50 | + | Get(key []byte) ([]byte, bool) | |
| 50 | 51 | Set(key, val []byte) error | |
| 51 | 52 | Delete(key []byte) error | |
| 52 | 53 | Range(fn func(k, v []byte) bool) | |
| @@ -47,6 +47,7 @@ type Listener interface { | |||||
| 47 | // | 47 | // | ||
| 48 | // Implementations may buffer writes; Sync is the durability barrier. | 48 | // Implementations may buffer writes; Sync is the durability barrier. | ||
| 49 | type Store interface { | 49 | type Store interface { | ||
| 50 | + | Get(key []byte) ([]byte, bool) | |||
| 50 | Set(key, val []byte) error | 51 | Set(key, val []byte) error | ||
| 51 | Delete(key []byte) error | 52 | Delete(key []byte) error | ||
| 52 | Range(fn func(k, v []byte) bool) | 53 | Range(fn func(k, v []byte) bool) | ||
M internal/usecase/daemon/client.go +9 -3
| @@ -366,7 +366,10 @@ func (d *Daemon) clientGone(sess *session, ac *attachedClient, explicit bool) { | |||
| 366 | 366 | return // already displaced by a newer client; nothing to do | |
| 367 | 367 | } | |
| 368 | 368 | d.unregisterPreview(ac) | |
| 369 | - | if !sess.ephemeral { | |
| 369 | + | sess.mu.Lock() | |
| 370 | + | ephemeral := sess.ephemeral | |
| 371 | + | sess.mu.Unlock() | |
| 372 | + | if !ephemeral { | |
| 370 | 373 | d.refreshSessionCwd(sess) | |
| 371 | 374 | } | |
| 372 | 375 | d.resetScreenDefaultColors(sess) | |
| @@ -378,7 +381,7 @@ func (d *Daemon) clientGone(sess *session, ac *attachedClient, explicit bool) { | |||
| 378 | 381 | d.boundedSend(ac, frameDetached(ports.ReasonDetach)) | |
| 379 | 382 | } | |
| 380 | 383 | _ = ac.tr.Close() | |
| 381 | - | if sess.ephemeral { | |
| 384 | + | if ephemeral { | |
| 382 | 385 | _ = d.killSession(sess, ports.ReasonSessionKilled, false) | |
| 383 | 386 | } | |
| 384 | 387 | } | |
| @@ -448,7 +451,10 @@ func (d *Daemon) detachOnSendError(sess *session, ac *attachedClient) { | |||
| 448 | 451 | d.resetScreenDefaultColors(sess) | |
| 449 | 452 | _ = ac.tr.Close() | |
| 450 | 453 | d.log.Warn("detached client after send error", "session", sess.name) | |
| 451 | - | if sess.ephemeral { | |
| 454 | + | sess.mu.Lock() | |
| 455 | + | ephemeral := sess.ephemeral | |
| 456 | + | sess.mu.Unlock() | |
| 457 | + | if ephemeral { | |
| 452 | 458 | _ = d.killSession(sess, ports.ReasonSessionKilled, false) | |
| 453 | 459 | } | |
| 454 | 460 | } | |
| @@ -366,7 +366,10 @@ func (d *Daemon) clientGone(sess *session, ac *attachedClient, explicit bool) { | |||||
| 366 | return // already displaced by a newer client; nothing to do | 366 | return // already displaced by a newer client; nothing to do | ||
| 367 | } | 367 | } | ||
| 368 | d.unregisterPreview(ac) | 368 | d.unregisterPreview(ac) | ||
| 369 | - | if !sess.ephemeral { | 369 | + | sess.mu.Lock() |
| 370 | + | ephemeral := sess.ephemeral | |||
| 371 | + | sess.mu.Unlock() | |||
| 372 | + | if !ephemeral { | |||
| 370 | d.refreshSessionCwd(sess) | 373 | d.refreshSessionCwd(sess) | ||
| 371 | } | 374 | } | ||
| 372 | d.resetScreenDefaultColors(sess) | 375 | d.resetScreenDefaultColors(sess) | ||
| @@ -378,7 +381,7 @@ func (d *Daemon) clientGone(sess *session, ac *attachedClient, explicit bool) { | |||||
| 378 | d.boundedSend(ac, frameDetached(ports.ReasonDetach)) | 381 | d.boundedSend(ac, frameDetached(ports.ReasonDetach)) | ||
| 379 | } | 382 | } | ||
| 380 | _ = ac.tr.Close() | 383 | _ = ac.tr.Close() | ||
| 381 | - | if sess.ephemeral { | 384 | + | if ephemeral { |
| 382 | _ = d.killSession(sess, ports.ReasonSessionKilled, false) | 385 | _ = d.killSession(sess, ports.ReasonSessionKilled, false) | ||
| 383 | } | 386 | } | ||
| 384 | } | 387 | } | ||
| @@ -448,7 +451,10 @@ func (d *Daemon) detachOnSendError(sess *session, ac *attachedClient) { | |||||
| 448 | d.resetScreenDefaultColors(sess) | 451 | d.resetScreenDefaultColors(sess) | ||
| 449 | _ = ac.tr.Close() | 452 | _ = ac.tr.Close() | ||
| 450 | d.log.Warn("detached client after send error", "session", sess.name) | 453 | d.log.Warn("detached client after send error", "session", sess.name) | ||
| 451 | - | if sess.ephemeral { | 454 | + | sess.mu.Lock() |
| 455 | + | ephemeral := sess.ephemeral | |||
| 456 | + | sess.mu.Unlock() | |||
| 457 | + | if ephemeral { | |||
| 452 | _ = d.killSession(sess, ports.ReasonSessionKilled, false) | 458 | _ = d.killSession(sess, ports.ReasonSessionKilled, false) | ||
| 453 | } | 459 | } | ||
| 454 | } | 460 | } | ||
M internal/usecase/daemon/daemon.go +12 -7
| @@ -120,6 +120,7 @@ type stoppedSession struct { | |||
| 120 | 120 | name string | |
| 121 | 121 | cwd string | |
| 122 | 122 | createdAt int64 | |
| 123 | + | purging bool | |
| 123 | 124 | } | |
| 124 | 125 | ||
| 125 | 126 | type Option func(*Daemon) | |
| @@ -206,9 +207,11 @@ func (d *Daemon) Serve(ctx context.Context, l ports.Listener) error { | |||
| 206 | 207 | d.sessWg.Go(func() { | |
| 207 | 208 | d.attentionAnimator(d.serveCtx) | |
| 208 | 209 | }) | |
| 209 | - | d.sessWg.Go(func() { | |
| 210 | - | d.cwdSampler(d.serveCtx) | |
| 211 | - | }) | |
| 210 | + | if d.persistEnabled && d.procCwd != nil { | |
| 211 | + | d.sessWg.Go(func() { | |
| 212 | + | d.cwdSampler(d.serveCtx) | |
| 213 | + | }) | |
| 214 | + | } | |
| 212 | 215 | ||
| 213 | 216 | // Break the accept loop when either the parent context is cancelled or the | |
| 214 | 217 | // registry drains to empty: both close the listener, which fails Accept. | |
| @@ -326,7 +329,10 @@ func (d *Daemon) handleList(tr ports.Transport) { | |||
| 326 | 329 | s.mu.Unlock() | |
| 327 | 330 | infos = append(infos, info) | |
| 328 | 331 | } | |
| 329 | - | for name := range d.stopped { | |
| 332 | + | for name, stopped := range d.stopped { | |
| 333 | + | if stopped.purging { | |
| 334 | + | continue | |
| 335 | + | } | |
| 330 | 336 | if _, live := liveNames[name]; live { | |
| 331 | 337 | continue | |
| 332 | 338 | } | |
| @@ -358,13 +364,12 @@ func (d *Daemon) handleKill(tr ports.Transport, f ports.Frame) { | |||
| 358 | 364 | target := d.findByNameLocked(k.Name) | |
| 359 | 365 | if target == nil { | |
| 360 | 366 | if stopped, ok := d.stopped[k.Name]; ok { | |
| 361 | - | d.mu.Unlock() | |
| 362 | 367 | if err := d.persist.Delete(k.Name); err != nil { | |
| 368 | + | d.mu.Unlock() | |
| 363 | 369 | d.log.Warn("deleting persisted stopped session failed", "err", err, "session", k.Name) | |
| 364 | 370 | _ = tr.Send(frameError(ports.ErrInternal, "deleting persisted stopped session failed")) | |
| 365 | 371 | return | |
| 366 | 372 | } | |
| 367 | - | d.mu.Lock() | |
| 368 | 373 | if cur, ok := d.stopped[k.Name]; ok && cur == stopped { | |
| 369 | 374 | delete(d.stopped, k.Name) | |
| 370 | 375 | } | |
| @@ -483,7 +488,7 @@ func (d *Daemon) route(h ports.Hello, tr ports.Transport) (*session, *attachedCl | |||
| 483 | 488 | sess := d.findByNameLocked(h.Name) | |
| 484 | 489 | if sess == nil { | |
| 485 | 490 | stopped, ok := d.stopped[h.Name] | |
| 486 | - | if !ok { | |
| 491 | + | if !ok || stopped.purging { | |
| 487 | 492 | d.mu.Unlock() | |
| 488 | 493 | return nil, nil, &protoErr{ports.ErrNoSuchSession, "no such session: " + h.Name} | |
| 489 | 494 | } | |
| @@ -120,6 +120,7 @@ type stoppedSession struct { | |||||
| 120 | name string | 120 | name string | ||
| 121 | cwd string | 121 | cwd string | ||
| 122 | createdAt int64 | 122 | createdAt int64 | ||
| 123 | + | purging bool | |||
| 123 | } | 124 | } | ||
| 124 | 125 | ||||
| 125 | type Option func(*Daemon) | 126 | type Option func(*Daemon) | ||
| @@ -206,9 +207,11 @@ func (d *Daemon) Serve(ctx context.Context, l ports.Listener) error { | |||||
| 206 | d.sessWg.Go(func() { | 207 | d.sessWg.Go(func() { | ||
| 207 | d.attentionAnimator(d.serveCtx) | 208 | d.attentionAnimator(d.serveCtx) | ||
| 208 | }) | 209 | }) | ||
| 209 | - | d.sessWg.Go(func() { | 210 | + | if d.persistEnabled && d.procCwd != nil { |
| 210 | - | d.cwdSampler(d.serveCtx) | 211 | + | d.sessWg.Go(func() { |
| 211 | - | }) | 212 | + | d.cwdSampler(d.serveCtx) |
| 213 | + | }) | |||
| 214 | + | } | |||
| 212 | 215 | ||||
| 213 | // Break the accept loop when either the parent context is cancelled or the | 216 | // Break the accept loop when either the parent context is cancelled or the | ||
| 214 | // registry drains to empty: both close the listener, which fails Accept. | 217 | // registry drains to empty: both close the listener, which fails Accept. | ||
| @@ -326,7 +329,10 @@ func (d *Daemon) handleList(tr ports.Transport) { | |||||
| 326 | s.mu.Unlock() | 329 | s.mu.Unlock() | ||
| 327 | infos = append(infos, info) | 330 | infos = append(infos, info) | ||
| 328 | } | 331 | } | ||
| 329 | - | for name := range d.stopped { | 332 | + | for name, stopped := range d.stopped { |
| 333 | + | if stopped.purging { | |||
| 334 | + | continue | |||
| 335 | + | } | |||
| 330 | if _, live := liveNames[name]; live { | 336 | if _, live := liveNames[name]; live { | ||
| 331 | continue | 337 | continue | ||
| 332 | } | 338 | } | ||
| @@ -358,13 +364,12 @@ func (d *Daemon) handleKill(tr ports.Transport, f ports.Frame) { | |||||
| 358 | target := d.findByNameLocked(k.Name) | 364 | target := d.findByNameLocked(k.Name) | ||
| 359 | if target == nil { | 365 | if target == nil { | ||
| 360 | if stopped, ok := d.stopped[k.Name]; ok { | 366 | if stopped, ok := d.stopped[k.Name]; ok { | ||
| 361 | - | d.mu.Unlock() | |||
| 362 | if err := d.persist.Delete(k.Name); err != nil { | 367 | if err := d.persist.Delete(k.Name); err != nil { | ||
| 368 | + | d.mu.Unlock() | |||
| 363 | d.log.Warn("deleting persisted stopped session failed", "err", err, "session", k.Name) | 369 | d.log.Warn("deleting persisted stopped session failed", "err", err, "session", k.Name) | ||
| 364 | _ = tr.Send(frameError(ports.ErrInternal, "deleting persisted stopped session failed")) | 370 | _ = tr.Send(frameError(ports.ErrInternal, "deleting persisted stopped session failed")) | ||
| 365 | return | 371 | return | ||
| 366 | } | 372 | } | ||
| 367 | - | d.mu.Lock() | |||
| 368 | if cur, ok := d.stopped[k.Name]; ok && cur == stopped { | 373 | if cur, ok := d.stopped[k.Name]; ok && cur == stopped { | ||
| 369 | delete(d.stopped, k.Name) | 374 | delete(d.stopped, k.Name) | ||
| 370 | } | 375 | } | ||
| @@ -483,7 +488,7 @@ func (d *Daemon) route(h ports.Hello, tr ports.Transport) (*session, *attachedCl | |||||
| 483 | sess := d.findByNameLocked(h.Name) | 488 | sess := d.findByNameLocked(h.Name) | ||
| 484 | if sess == nil { | 489 | if sess == nil { | ||
| 485 | stopped, ok := d.stopped[h.Name] | 490 | stopped, ok := d.stopped[h.Name] | ||
| 486 | - | if !ok { | 491 | + | if !ok || stopped.purging { |
| 487 | d.mu.Unlock() | 492 | d.mu.Unlock() | ||
| 488 | return nil, nil, &protoErr{ports.ErrNoSuchSession, "no such session: " + h.Name} | 493 | return nil, nil, &protoErr{ports.ErrNoSuchSession, "no such session: " + h.Name} | ||
| 489 | } | 494 | } | ||
M internal/usecase/daemon/picker.go +7 -6
| @@ -56,7 +56,9 @@ func (d *Daemon) pickerViews(cur *session) ([]picker.SessionView, int) { | |||
| 56 | 56 | sessions := d.sessionsSnapshotLocked() | |
| 57 | 57 | stopped := make([]stoppedSession, 0, len(d.stopped)) | |
| 58 | 58 | for _, s := range d.stopped { | |
| 59 | - | stopped = append(stopped, s) | |
| 59 | + | if !s.purging { | |
| 60 | + | stopped = append(stopped, s) | |
| 61 | + | } | |
| 60 | 62 | } | |
| 61 | 63 | d.mu.Unlock() | |
| 62 | 64 | sort.Slice(sessions, func(i, j int) bool { return sessions[i].name < sessions[j].name }) | |
| @@ -406,7 +408,7 @@ func (d *Daemon) stealClientForTarget(from *session, ac *attachedClient, targetS | |||
| 406 | 408 | func (d *Daemon) resumeStoppedAndSwitch(from *session, ac *attachedClient, target picker.Target) { | |
| 407 | 409 | d.mu.Lock() | |
| 408 | 410 | stopped, ok := d.stopped[target.Name] | |
| 409 | - | if !ok { | |
| 411 | + | if !ok || stopped.purging { | |
| 410 | 412 | d.mu.Unlock() | |
| 411 | 413 | d.paint(from, ac, true) | |
| 412 | 414 | return | |
| @@ -441,18 +443,17 @@ func (d *Daemon) killPickerTarget(target picker.Target) { | |||
| 441 | 443 | if target.Stopped { | |
| 442 | 444 | d.mu.Lock() | |
| 443 | 445 | stopped, ok := d.stopped[target.Name] | |
| 444 | - | d.mu.Unlock() | |
| 445 | - | if ok { | |
| 446 | + | if ok && !stopped.purging { | |
| 446 | 447 | if err := d.persist.Delete(target.Name); err != nil { | |
| 448 | + | d.mu.Unlock() | |
| 447 | 449 | d.log.Warn("deleting persisted stopped session failed", "err", err, "session", target.Name) | |
| 448 | 450 | return | |
| 449 | 451 | } | |
| 450 | - | d.mu.Lock() | |
| 451 | 452 | if cur, ok := d.stopped[target.Name]; ok && cur == stopped { | |
| 452 | 453 | delete(d.stopped, target.Name) | |
| 453 | 454 | } | |
| 454 | - | d.mu.Unlock() | |
| 455 | 455 | } | |
| 456 | + | d.mu.Unlock() | |
| 456 | 457 | return | |
| 457 | 458 | } | |
| 458 | 459 | d.mu.Lock() | |
| @@ -56,7 +56,9 @@ func (d *Daemon) pickerViews(cur *session) ([]picker.SessionView, int) { | |||||
| 56 | sessions := d.sessionsSnapshotLocked() | 56 | sessions := d.sessionsSnapshotLocked() | ||
| 57 | stopped := make([]stoppedSession, 0, len(d.stopped)) | 57 | stopped := make([]stoppedSession, 0, len(d.stopped)) | ||
| 58 | for _, s := range d.stopped { | 58 | for _, s := range d.stopped { | ||
| 59 | - | stopped = append(stopped, s) | 59 | + | if !s.purging { |
| 60 | + | stopped = append(stopped, s) | |||
| 61 | + | } | |||
| 60 | } | 62 | } | ||
| 61 | d.mu.Unlock() | 63 | d.mu.Unlock() | ||
| 62 | sort.Slice(sessions, func(i, j int) bool { return sessions[i].name < sessions[j].name }) | 64 | sort.Slice(sessions, func(i, j int) bool { return sessions[i].name < sessions[j].name }) | ||
| @@ -406,7 +408,7 @@ func (d *Daemon) stealClientForTarget(from *session, ac *attachedClient, targetS | |||||
| 406 | func (d *Daemon) resumeStoppedAndSwitch(from *session, ac *attachedClient, target picker.Target) { | 408 | func (d *Daemon) resumeStoppedAndSwitch(from *session, ac *attachedClient, target picker.Target) { | ||
| 407 | d.mu.Lock() | 409 | d.mu.Lock() | ||
| 408 | stopped, ok := d.stopped[target.Name] | 410 | stopped, ok := d.stopped[target.Name] | ||
| 409 | - | if !ok { | 411 | + | if !ok || stopped.purging { |
| 410 | d.mu.Unlock() | 412 | d.mu.Unlock() | ||
| 411 | d.paint(from, ac, true) | 413 | d.paint(from, ac, true) | ||
| 412 | return | 414 | return | ||
| @@ -441,18 +443,17 @@ func (d *Daemon) killPickerTarget(target picker.Target) { | |||||
| 441 | if target.Stopped { | 443 | if target.Stopped { | ||
| 442 | d.mu.Lock() | 444 | d.mu.Lock() | ||
| 443 | stopped, ok := d.stopped[target.Name] | 445 | stopped, ok := d.stopped[target.Name] | ||
| 444 | - | d.mu.Unlock() | 446 | + | if ok && !stopped.purging { |
| 445 | - | if ok { | |||
| 446 | if err := d.persist.Delete(target.Name); err != nil { | 447 | if err := d.persist.Delete(target.Name); err != nil { | ||
| 448 | + | d.mu.Unlock() | |||
| 447 | d.log.Warn("deleting persisted stopped session failed", "err", err, "session", target.Name) | 449 | d.log.Warn("deleting persisted stopped session failed", "err", err, "session", target.Name) | ||
| 448 | return | 450 | return | ||
| 449 | } | 451 | } | ||
| 450 | - | d.mu.Lock() | |||
| 451 | if cur, ok := d.stopped[target.Name]; ok && cur == stopped { | 452 | if cur, ok := d.stopped[target.Name]; ok && cur == stopped { | ||
| 452 | delete(d.stopped, target.Name) | 453 | delete(d.stopped, target.Name) | ||
| 453 | } | 454 | } | ||
| 454 | - | d.mu.Unlock() | |||
| 455 | } | 455 | } | ||
| 456 | + | d.mu.Unlock() | |||
| 456 | return | 457 | return | ||
| 457 | } | 458 | } | ||
| 458 | d.mu.Lock() | 459 | d.mu.Lock() | ||
M internal/usecase/daemon/session.go +18 -11
| @@ -293,13 +293,16 @@ func (d *Daemon) renameSession(sess *session, name string) error { | |||
| 293 | 293 | createdAt = time.Now().UnixNano() | |
| 294 | 294 | sess.createdAt = createdAt | |
| 295 | 295 | } | |
| 296 | - | if !wasEphemeral && oldName != name { | |
| 297 | - | if err := d.persist.Delete(oldName); err != nil { | |
| 296 | + | if wasEphemeral || oldName != name { | |
| 297 | + | if err := d.persist.Save(persist.Record{Name: name, Cwd: sess.cwd, CreatedAt: createdAt, UpdatedAt: time.Now().UnixNano()}); err != nil { | |
| 298 | 298 | return err | |
| 299 | 299 | } | |
| 300 | 300 | } | |
| 301 | - | if wasEphemeral || oldName != name { | |
| 302 | - | if err := d.persist.Save(persist.Record{Name: name, Cwd: sess.cwd, CreatedAt: createdAt, UpdatedAt: time.Now().UnixNano()}); err != nil { | |
| 301 | + | if !wasEphemeral && oldName != name { | |
| 302 | + | if err := d.persist.Delete(oldName); err != nil { | |
| 303 | + | if cleanupErr := d.persist.Delete(name); cleanupErr != nil { | |
| 304 | + | d.log.Warn("cleaning up renamed persisted session failed", "err", cleanupErr, "session", name) | |
| 305 | + | } | |
| 303 | 306 | return err | |
| 304 | 307 | } | |
| 305 | 308 | } | |
| @@ -375,11 +378,8 @@ func (d *Daemon) killSession(sess *session, reason uint8, purge bool) error { | |||
| 375 | 378 | ephemeral := sess.ephemeral | |
| 376 | 379 | sess.mu.Unlock() | |
| 377 | 380 | if !ephemeral { | |
| 378 | - | if purge { | |
| 379 | - | delete(d.stopped, stoppedName) | |
| 380 | - | } else { | |
| 381 | - | d.stopped[stoppedName] = stoppedSession{name: stoppedName, cwd: stoppedCwd, createdAt: createdAt} | |
| 382 | - | } | |
| 381 | + | stopped := stoppedSession{name: stoppedName, cwd: stoppedCwd, createdAt: createdAt, purging: purge} | |
| 382 | + | d.stopped[stoppedName] = stopped | |
| 383 | 383 | } | |
| 384 | 384 | empty := len(d.sessions) == 0 | |
| 385 | 385 | if empty { | |
| @@ -396,8 +396,15 @@ func (d *Daemon) killSession(sess *session, reason uint8, purge bool) error { | |||
| 396 | 396 | purgeErr = err | |
| 397 | 397 | d.log.Warn("deleting persisted session failed", "err", err, "session", stoppedName) | |
| 398 | 398 | d.mu.Lock() | |
| 399 | - | if _, live := d.sessions[sess.id]; !live { | |
| 400 | - | d.stopped[stoppedName] = stoppedSession{name: stoppedName, cwd: stoppedCwd, createdAt: createdAt} | |
| 399 | + | if stopped, ok := d.stopped[stoppedName]; ok && stopped.purging { | |
| 400 | + | stopped.purging = false | |
| 401 | + | d.stopped[stoppedName] = stopped | |
| 402 | + | } | |
| 403 | + | d.mu.Unlock() | |
| 404 | + | } else { | |
| 405 | + | d.mu.Lock() | |
| 406 | + | if stopped, ok := d.stopped[stoppedName]; ok && stopped.purging { | |
| 407 | + | delete(d.stopped, stoppedName) | |
| 401 | 408 | } | |
| 402 | 409 | d.mu.Unlock() | |
| 403 | 410 | } | |
| @@ -293,13 +293,16 @@ func (d *Daemon) renameSession(sess *session, name string) error { | |||||
| 293 | createdAt = time.Now().UnixNano() | 293 | createdAt = time.Now().UnixNano() | ||
| 294 | sess.createdAt = createdAt | 294 | sess.createdAt = createdAt | ||
| 295 | } | 295 | } | ||
| 296 | - | if !wasEphemeral && oldName != name { | 296 | + | if wasEphemeral || oldName != name { |
| 297 | - | if err := d.persist.Delete(oldName); err != nil { | 297 | + | if err := d.persist.Save(persist.Record{Name: name, Cwd: sess.cwd, CreatedAt: createdAt, UpdatedAt: time.Now().UnixNano()}); err != nil { |
| 298 | return err | 298 | return err | ||
| 299 | } | 299 | } | ||
| 300 | } | 300 | } | ||
| 301 | - | if wasEphemeral || oldName != name { | 301 | + | if !wasEphemeral && oldName != name { |
| 302 | - | if err := d.persist.Save(persist.Record{Name: name, Cwd: sess.cwd, CreatedAt: createdAt, UpdatedAt: time.Now().UnixNano()}); err != nil { | 302 | + | if err := d.persist.Delete(oldName); err != nil { |
| 303 | + | if cleanupErr := d.persist.Delete(name); cleanupErr != nil { | |||
| 304 | + | d.log.Warn("cleaning up renamed persisted session failed", "err", cleanupErr, "session", name) | |||
| 305 | + | } | |||
| 303 | return err | 306 | return err | ||
| 304 | } | 307 | } | ||
| 305 | } | 308 | } | ||
| @@ -375,11 +378,8 @@ func (d *Daemon) killSession(sess *session, reason uint8, purge bool) error { | |||||
| 375 | ephemeral := sess.ephemeral | 378 | ephemeral := sess.ephemeral | ||
| 376 | sess.mu.Unlock() | 379 | sess.mu.Unlock() | ||
| 377 | if !ephemeral { | 380 | if !ephemeral { | ||
| 378 | - | if purge { | 381 | + | stopped := stoppedSession{name: stoppedName, cwd: stoppedCwd, createdAt: createdAt, purging: purge} |
| 379 | - | delete(d.stopped, stoppedName) | 382 | + | d.stopped[stoppedName] = stopped |
| 380 | - | } else { | |||
| 381 | - | d.stopped[stoppedName] = stoppedSession{name: stoppedName, cwd: stoppedCwd, createdAt: createdAt} | |||
| 382 | - | } | |||
| 383 | } | 383 | } | ||
| 384 | empty := len(d.sessions) == 0 | 384 | empty := len(d.sessions) == 0 | ||
| 385 | if empty { | 385 | if empty { | ||
| @@ -396,8 +396,15 @@ func (d *Daemon) killSession(sess *session, reason uint8, purge bool) error { | |||||
| 396 | purgeErr = err | 396 | purgeErr = err | ||
| 397 | d.log.Warn("deleting persisted session failed", "err", err, "session", stoppedName) | 397 | d.log.Warn("deleting persisted session failed", "err", err, "session", stoppedName) | ||
| 398 | d.mu.Lock() | 398 | d.mu.Lock() | ||
| 399 | - | if _, live := d.sessions[sess.id]; !live { | 399 | + | if stopped, ok := d.stopped[stoppedName]; ok && stopped.purging { |
| 400 | - | d.stopped[stoppedName] = stoppedSession{name: stoppedName, cwd: stoppedCwd, createdAt: createdAt} | 400 | + | stopped.purging = false |
| 401 | + | d.stopped[stoppedName] = stopped | |||
| 402 | + | } | |||
| 403 | + | d.mu.Unlock() | |||
| 404 | + | } else { | |||
| 405 | + | d.mu.Lock() | |||
| 406 | + | if stopped, ok := d.stopped[stoppedName]; ok && stopped.purging { | |||
| 407 | + | delete(d.stopped, stoppedName) | |||
| 401 | } | 408 | } | ||
| 402 | d.mu.Unlock() | 409 | d.mu.Unlock() | ||
| 403 | } | 410 | } | ||
M internal/usecase/daemon/session_test.go +6 -0
| @@ -596,6 +596,12 @@ func newMockStore(t *testing.T) (*portsmocks.MockStore, *mockStoreState) { | |||
| 596 | 596 | t.Helper() | |
| 597 | 597 | state := &mockStoreState{data: make(map[string][]byte)} | |
| 598 | 598 | store := portsmocks.NewMockStore(t) | |
| 599 | + | store.EXPECT().Get(mock.Anything).RunAndReturn(func(k []byte) ([]byte, bool) { | |
| 600 | + | state.mu.Lock() | |
| 601 | + | defer state.mu.Unlock() | |
| 602 | + | v, ok := state.data[string(k)] | |
| 603 | + | return append([]byte(nil), v...), ok | |
| 604 | + | }).Maybe() | |
| 599 | 605 | store.EXPECT().Set(mock.Anything, mock.Anything).RunAndReturn(func(k, v []byte) error { | |
| 600 | 606 | state.mu.Lock() | |
| 601 | 607 | defer state.mu.Unlock() | |
| @@ -596,6 +596,12 @@ func newMockStore(t *testing.T) (*portsmocks.MockStore, *mockStoreState) { | |||||
| 596 | t.Helper() | 596 | t.Helper() | ||
| 597 | state := &mockStoreState{data: make(map[string][]byte)} | 597 | state := &mockStoreState{data: make(map[string][]byte)} | ||
| 598 | store := portsmocks.NewMockStore(t) | 598 | store := portsmocks.NewMockStore(t) | ||
| 599 | + | store.EXPECT().Get(mock.Anything).RunAndReturn(func(k []byte) ([]byte, bool) { | |||
| 600 | + | state.mu.Lock() | |||
| 601 | + | defer state.mu.Unlock() | |||
| 602 | + | v, ok := state.data[string(k)] | |||
| 603 | + | return append([]byte(nil), v...), ok | |||
| 604 | + | }).Maybe() | |||
| 599 | store.EXPECT().Set(mock.Anything, mock.Anything).RunAndReturn(func(k, v []byte) error { | 605 | store.EXPECT().Set(mock.Anything, mock.Anything).RunAndReturn(func(k, v []byte) error { | ||
| 600 | state.mu.Lock() | 606 | state.mu.Lock() | ||
| 601 | defer state.mu.Unlock() | 607 | defer state.mu.Unlock() | ||
M pkg/kv/codec.go +4 -0
| @@ -4,6 +4,7 @@ import ( | |||
| 4 | 4 | "encoding/binary" | |
| 5 | 5 | "errors" | |
| 6 | 6 | "hash/crc32" | |
| 7 | + | "math" | |
| 7 | 8 | ) | |
| 8 | 9 | ||
| 9 | 10 | const ( | |
| @@ -37,6 +38,9 @@ func encodeRecord(op byte, key, value []byte) ([]byte, error) { | |||
| 37 | 38 | if op == opSet { | |
| 38 | 39 | payloadLen += len(value) | |
| 39 | 40 | } | |
| 41 | + | if uint64(payloadLen) > math.MaxUint32 { | |
| 42 | + | return nil, errors.New("payload too large") | |
| 43 | + | } | |
| 40 | 44 | ||
| 41 | 45 | buf := make([]byte, headerLen+payloadLen) | |
| 42 | 46 | payload := buf[headerLen:] | |
| @@ -4,6 +4,7 @@ import ( | |||||
| 4 | "encoding/binary" | 4 | "encoding/binary" | ||
| 5 | "errors" | 5 | "errors" | ||
| 6 | "hash/crc32" | 6 | "hash/crc32" | ||
| 7 | + | "math" | |||
| 7 | ) | 8 | ) | ||
| 8 | 9 | ||||
| 9 | const ( | 10 | const ( | ||
| @@ -37,6 +38,9 @@ func encodeRecord(op byte, key, value []byte) ([]byte, error) { | |||||
| 37 | if op == opSet { | 38 | if op == opSet { | ||
| 38 | payloadLen += len(value) | 39 | payloadLen += len(value) | ||
| 39 | } | 40 | } | ||
| 41 | + | if uint64(payloadLen) > math.MaxUint32 { | |||
| 42 | + | return nil, errors.New("payload too large") | |||
| 43 | + | } | |||
| 40 | 44 | ||||
| 41 | buf := make([]byte, headerLen+payloadLen) | 45 | buf := make([]byte, headerLen+payloadLen) | ||
| 42 | payload := buf[headerLen:] | 46 | payload := buf[headerLen:] | ||
M pkg/kv/kv.go +18 -3
| @@ -201,10 +201,14 @@ func (s *Store) appendRecord(op byte, key, val []byte) error { | |||
| 201 | 201 | if err != nil { | |
| 202 | 202 | return err | |
| 203 | 203 | } | |
| 204 | - | if _, err := s.file.Write(buf); err != nil { | |
| 204 | + | n, err := s.file.Write(buf) | |
| 205 | + | if err != nil || n != len(buf) { | |
| 205 | 206 | _ = s.file.Truncate(pos) | |
| 206 | 207 | _, _ = s.file.Seek(pos, io.SeekStart) | |
| 207 | - | return err | |
| 208 | + | if err != nil { | |
| 209 | + | return err | |
| 210 | + | } | |
| 211 | + | return io.ErrShortWrite | |
| 208 | 212 | } | |
| 209 | 213 | ||
| 210 | 214 | k := string(key) | |
| @@ -271,25 +275,36 @@ func (s *Store) compactLocked() error { | |||
| 271 | 275 | _ = os.Remove(tmp) | |
| 272 | 276 | return err | |
| 273 | 277 | } | |
| 274 | - | if err := s.file.Close(); err != nil { | |
| 278 | + | bak := s.path + ".bak" | |
| 279 | + | _ = os.Remove(bak) | |
| 280 | + | if err := os.Rename(s.path, bak); err != nil { | |
| 275 | 281 | _ = os.Remove(tmp) | |
| 276 | 282 | return err | |
| 277 | 283 | } | |
| 278 | 284 | if err := os.Rename(tmp, s.path); err != nil { | |
| 285 | + | _ = os.Rename(bak, s.path) | |
| 286 | + | _ = os.Remove(tmp) | |
| 279 | 287 | return err | |
| 280 | 288 | } | |
| 281 | 289 | f, err := os.OpenFile(s.path, os.O_RDWR, 0o600) | |
| 282 | 290 | if err != nil { | |
| 291 | + | _ = os.Remove(s.path) | |
| 292 | + | _ = os.Rename(bak, s.path) | |
| 283 | 293 | return err | |
| 284 | 294 | } | |
| 285 | 295 | if _, err := f.Seek(0, io.SeekEnd); err != nil { | |
| 286 | 296 | _ = f.Close() | |
| 297 | + | _ = os.Remove(s.path) | |
| 298 | + | _ = os.Rename(bak, s.path) | |
| 287 | 299 | return err | |
| 288 | 300 | } | |
| 301 | + | old := s.file | |
| 289 | 302 | s.file = f | |
| 290 | 303 | s.entrySize = newSizes | |
| 291 | 304 | s.total = newTotal | |
| 292 | 305 | s.live = newLive | |
| 306 | + | _ = old.Close() | |
| 307 | + | _ = os.Remove(bak) | |
| 293 | 308 | return nil | |
| 294 | 309 | } | |
| 295 | 310 | ||
| @@ -201,10 +201,14 @@ func (s *Store) appendRecord(op byte, key, val []byte) error { | |||||
| 201 | if err != nil { | 201 | if err != nil { | ||
| 202 | return err | 202 | return err | ||
| 203 | } | 203 | } | ||
| 204 | - | if _, err := s.file.Write(buf); err != nil { | 204 | + | n, err := s.file.Write(buf) |
| 205 | + | if err != nil || n != len(buf) { | |||
| 205 | _ = s.file.Truncate(pos) | 206 | _ = s.file.Truncate(pos) | ||
| 206 | _, _ = s.file.Seek(pos, io.SeekStart) | 207 | _, _ = s.file.Seek(pos, io.SeekStart) | ||
| 207 | - | return err | 208 | + | if err != nil { |
| 209 | + | return err | |||
| 210 | + | } | |||
| 211 | + | return io.ErrShortWrite | |||
| 208 | } | 212 | } | ||
| 209 | 213 | ||||
| 210 | k := string(key) | 214 | k := string(key) | ||
| @@ -271,25 +275,36 @@ func (s *Store) compactLocked() error { | |||||
| 271 | _ = os.Remove(tmp) | 275 | _ = os.Remove(tmp) | ||
| 272 | return err | 276 | return err | ||
| 273 | } | 277 | } | ||
| 274 | - | if err := s.file.Close(); err != nil { | 278 | + | bak := s.path + ".bak" |
| 279 | + | _ = os.Remove(bak) | |||
| 280 | + | if err := os.Rename(s.path, bak); err != nil { | |||
| 275 | _ = os.Remove(tmp) | 281 | _ = os.Remove(tmp) | ||
| 276 | return err | 282 | return err | ||
| 277 | } | 283 | } | ||
| 278 | if err := os.Rename(tmp, s.path); err != nil { | 284 | if err := os.Rename(tmp, s.path); err != nil { | ||
| 285 | + | _ = os.Rename(bak, s.path) | |||
| 286 | + | _ = os.Remove(tmp) | |||
| 279 | return err | 287 | return err | ||
| 280 | } | 288 | } | ||
| 281 | f, err := os.OpenFile(s.path, os.O_RDWR, 0o600) | 289 | f, err := os.OpenFile(s.path, os.O_RDWR, 0o600) | ||
| 282 | if err != nil { | 290 | if err != nil { | ||
| 291 | + | _ = os.Remove(s.path) | |||
| 292 | + | _ = os.Rename(bak, s.path) | |||
| 283 | return err | 293 | return err | ||
| 284 | } | 294 | } | ||
| 285 | if _, err := f.Seek(0, io.SeekEnd); err != nil { | 295 | if _, err := f.Seek(0, io.SeekEnd); err != nil { | ||
| 286 | _ = f.Close() | 296 | _ = f.Close() | ||
| 297 | + | _ = os.Remove(s.path) | |||
| 298 | + | _ = os.Rename(bak, s.path) | |||
| 287 | return err | 299 | return err | ||
| 288 | } | 300 | } | ||
| 301 | + | old := s.file | |||
| 289 | s.file = f | 302 | s.file = f | ||
| 290 | s.entrySize = newSizes | 303 | s.entrySize = newSizes | ||
| 291 | s.total = newTotal | 304 | s.total = newTotal | ||
| 292 | s.live = newLive | 305 | s.live = newLive | ||
| 306 | + | _ = old.Close() | |||
| 307 | + | _ = os.Remove(bak) | |||
| 293 | return nil | 308 | return nil | ||
| 294 | } | 309 | } | ||
| 295 | 310 | ||||