1
package chord
2
 
3
import (
4
	"context"
5
	"fmt"
6
	"sync"
7
	"time"
8
 
9
	"go.miragespace.co/specter/spec/chord"
10
	"go.miragespace.co/specter/spec/protocol"
11
	"go.miragespace.co/specter/spec/rpc"
12
 
13
	"go.uber.org/zap"
14
	"golang.org/x/sync/errgroup"
15
)
16
 
17
// Workaround for https://github.com/golang/go/issues/49085#issuecomment-948108705
18
func kvMiddleware[V any](
19
	ctx context.Context,
20
	n *LocalNode,
21
	key []byte,
22
	handler func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (V, error),
23
) (V, error) {
24
	var zeroV V
25
	id := chord.Hash(key)
26
	reqCtx := rpc.GetContext(ctx)
27
	// if it is a replication request, bypass ownership checks
28
	if reqCtx.GetRequestTarget() == protocol.Context_KV_REPLICATION {
29
		return handler(ctx, n.kv, targetReplication, id)
30
	}
31
	// otherwise, continue with the usual lookup + forward/handle
32
	succ, err := n.FindSuccessor(id)
33
	switch err {
34
	case nil:
35
	case chord.ErrNodeGone:
36
		// if the remote node happens to be leaving, the caller needs to retry
37
		return zeroV, chord.ErrKVStaleOwnership
38
	default:
39
		return zeroV, err
40
	}
41
 
42
	if succ.ID() != n.ID() {
43
		// remote KV
44
		return handler(ctx, succ, targetRemote, id)
45
	}
46
 
47
	l := n.logger.With(
48
		zap.String("key", string(key)),
49
		zap.Uint64("id", id),
50
	)
51
 
52
	// local KV
53
	n.surrogateMu.RLock()
54
	defer n.surrogateMu.RUnlock()
55
 
56
	// maybe we are joining or leaving
57
	state := n.state.Get()
58
	if state != chord.Active {
59
		l.Debug("KV Handler node is not in active state", zap.String("state", state.String()))
60
		n.kvStaleCount.Inc()
61
		return zeroV, chord.ErrKVStaleOwnership
62
	}
63
 
64
	if n.surrogate != nil {
65
		l = l.With(zap.Object("surrogate", n.surrogate.Identity()))
66
	}
67
 
68
	n.predecessorMu.RLock()
69
	defer n.predecessorMu.RUnlock()
70
 
71
	if n.predecessor != nil {
72
		l = l.With(zap.Object("predecessor", n.predecessor.Identity()))
73
	}
74
 
75
	if n.surrogate != nil && chord.Between(n.ID(), id, n.surrogate.Identity().GetId(), true) {
76
		l.Warn("KV Ownership moved, forwarding to surrogate")
77
		n.kvStaleCount.Inc()
78
		return handler(ctx, n.surrogate, targetSurrogate, id)
79
	}
80
 
81
	if n.predecessor != nil && !chord.Between(n.predecessor.ID(), id, n.ID(), true) {
82
		l.Debug("Key not in range")
83
		n.kvStaleCount.Inc()
84
		return zeroV, chord.ErrKVStaleOwnership
85
	}
86
 
87
	return handler(ctx, n.kv, targetLocal, id)
88
}
89
 
90
func (n *LocalNode) Put(ctx context.Context, key, value []byte) error {
91
	_, err := kvMiddleware(ctx, n, key,
92
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (any, error) {
93
			n.logger.Debug("KV Put", zap.String("target", target.String()), zap.String("key", string(key)), zap.Uint64("id", id))
94
			return nil, kv.Put(ctx, key, value)
95
		})
96
	return err
97
}
98
 
99
func (n *LocalNode) Get(ctx context.Context, key []byte) ([]byte, error) {
100
	return kvMiddleware(ctx, n, key,
101
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) ([]byte, error) {
102
			n.logger.Debug("KV Get", zap.String("target", target.String()), zap.String("key", string(key)), zap.Uint64("id", id))
103
			return kv.Get(ctx, key)
104
		})
105
}
106
 
107
func (n *LocalNode) Delete(ctx context.Context, key []byte) error {
108
	_, err := kvMiddleware(ctx, n, key,
109
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (any, error) {
110
			n.logger.Debug("KV Delete", zap.String("target", target.String()), zap.String("key", string(key)), zap.Uint64("id", id))
111
			return nil, kv.Delete(ctx, key)
112
		})
113
	return err
114
}
115
 
116
func (n *LocalNode) PrefixAppend(ctx context.Context, prefix []byte, child []byte) error {
117
	_, err := kvMiddleware(ctx, n, prefix,
118
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (any, error) {
119
			n.logger.Debug("KV PrefixAppend", zap.String("target", target.String()), zap.String("prefix", string(prefix)), zap.Uint64("id", id))
120
			return nil, kv.PrefixAppend(ctx, prefix, child)
121
		})
122
	return err
123
}
124
 
125
func (n *LocalNode) PrefixList(ctx context.Context, prefix []byte) ([][]byte, error) {
126
	return kvMiddleware(ctx, n, prefix,
127
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) ([][]byte, error) {
128
			n.logger.Debug("KV PrefixList", zap.String("target", target.String()), zap.String("prefix", string(prefix)), zap.Uint64("id", id))
129
			return kv.PrefixList(ctx, prefix)
130
		})
131
}
132
 
133
func (n *LocalNode) PrefixContains(ctx context.Context, prefix []byte, child []byte) (bool, error) {
134
	return kvMiddleware(ctx, n, prefix,
135
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (bool, error) {
136
			n.logger.Debug("KV PrefixContains", zap.String("target", target.String()), zap.String("prefix", string(prefix)), zap.Uint64("id", id))
137
			return kv.PrefixContains(ctx, prefix, child)
138
		})
139
}
140
 
141
func (n *LocalNode) PrefixRemove(ctx context.Context, prefix []byte, child []byte) error {
142
	_, err := kvMiddleware(ctx, n, prefix,
143
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (any, error) {
144
			n.logger.Debug("KV PrefixRemove", zap.String("target", target.String()), zap.String("prefix", string(prefix)), zap.Uint64("id", id))
145
			return nil, kv.PrefixRemove(ctx, prefix, child)
146
		})
147
	return err
148
}
149
 
150
func (n *LocalNode) Acquire(ctx context.Context, lease []byte, ttl time.Duration) (uint64, error) {
151
	return kvMiddleware(ctx, n, lease,
152
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (uint64, error) {
153
			n.logger.Debug("KV Acquire", zap.String("target", target.String()), zap.String("lease", string(lease)), zap.Uint64("id", id))
154
			return kv.Acquire(ctx, lease, ttl)
155
		})
156
}
157
 
158
func (n *LocalNode) Renew(ctx context.Context, lease []byte, ttl time.Duration, prevToken uint64) (uint64, error) {
159
	return kvMiddleware(ctx, n, lease,
160
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (uint64, error) {
161
			n.logger.Debug("KV Renew", zap.String("target", target.String()), zap.String("lease", string(lease)), zap.Uint64("id", id))
162
			return kv.Renew(ctx, lease, ttl, prevToken)
163
		})
164
}
165
 
166
func (n *LocalNode) Release(ctx context.Context, lease []byte, token uint64) error {
167
	_, err := kvMiddleware(ctx, n, lease,
168
		func(ctx context.Context, kv chord.KV, target kvTargetType, id uint64) (any, error) {
169
			n.logger.Debug("KV Release", zap.String("target", target.String()), zap.String("lease", string(lease)), zap.Uint64("id", id))
170
			return nil, kv.Release(ctx, lease, token)
171
		})
172
	return err
173
}
174
 
175
func (n *LocalNode) Import(ctx context.Context, keys [][]byte, values []*protocol.KVTransfer) error {
176
	state := n.state.Get()
177
	switch state {
178
	case chord.Inactive, chord.Leaving, chord.Left:
179
		return chord.ErrNodeGone
180
	}
181
	n.surrogateMu.Lock()
182
	defer n.surrogateMu.Unlock()
183
 
184
	n.logger.Debug("KV Import", zap.Int("num_keys", len(keys)))
185
	return n.kv.Import(ctx, keys, values)
186
}
187
 
188
func (n *LocalNode) ListKeys(ctx context.Context, prefix []byte) ([]*protocol.KeyComposite, error) {
189
	reqCtx := rpc.GetContext(ctx)
190
	if reqCtx.GetRequestTarget() == protocol.Context_KV_DIRECT_TARGET {
191
		n.logger.Debug("KV Listkeys", zap.Stringer("target", targetLocal), zap.String("prefix", string(prefix)))
192
		return func() ([]*protocol.KeyComposite, error) {
193
			n.surrogateMu.RLock()
194
			defer n.surrogateMu.RUnlock()
195
 
196
			state := n.state.Get()
197
			if state != chord.Active {
198
				n.kvStaleCount.Inc()
199
				return nil, chord.ErrKVStaleOwnership
200
			}
201
 
202
			return n.kv.ListKeys(ctx, prefix)
203
		}()
204
	}
205
	n.logger.Debug("KV Listkeys", zap.Stringer("target", targetRemote), zap.String("prefix", string(prefix)))
206
 
207
	var (
208
		keys              = make([]*protocol.KeyComposite, 0)
209
		nodes             = make([]chord.VNode, 0)
210
		seen              = make(map[uint64]bool)
211
		next  chord.VNode = n
212
		err   error
213
	)
214
 
215
	// we need to walk the entire ring
216
	for {
217
		next, err = n.FindSuccessor(chord.ModuloSum(next.ID(), 1))
218
		if err != nil {
219
			return nil, err
220
		}
221
		if next == nil {
222
			return nil, chord.ErrNodeNoSuccessor
223
		}
224
		if next.ID() == n.ID() {
225
			nodes = append(nodes, n)
226
			break
227
		}
228
		if seen[next.ID()] {
229
			return nil, fmt.Errorf("ring is unstable")
230
		}
231
		nodes = append(nodes, next)
232
		seen[next.ID()] = true
233
	}
234
 
235
	ctx = rpc.WithContext(ctx, &protocol.Context{
236
		RequestTarget: protocol.Context_KV_DIRECT_TARGET,
237
	})
238
 
239
	g, listCtx := errgroup.WithContext(ctx)
240
	g.SetLimit(10)
241
 
242
	readWg := sync.WaitGroup{}
243
	readWg.Add(1)
244
	resultCh := make(chan []*protocol.KeyComposite, 10)
245
	gErr := make(chan error, 1)
246
 
247
	go func() {
248
		defer readWg.Done()
249
 
250
		for k := range resultCh {
251
			keys = append(keys, k...)
252
		}
253
	}()
254
 
255
	for _, node := range nodes {
256
		g.Go(func() error {
257
			k, err := node.ListKeys(listCtx, prefix)
258
			if err != nil {
259
				return err
260
			}
261
			resultCh <- k
262
			return nil
263
		})
264
	}
265
 
266
	go func() {
267
		err := g.Wait()
268
		close(resultCh)
269
		readWg.Wait()
270
		gErr <- err
271
	}()
272
 
273
	if err := <-gErr; err != nil {
274
		return nil, err
275
	}
276
 
277
	return keys, nil
278
}