1
package acme
2
 
3
import (
4
	"context"
5
	"fmt"
6
	"io/fs"
7
	"strings"
8
	"sync"
9
	"sync/atomic"
10
	"time"
11
 
12
	"go.miragespace.co/specter/spec/chord"
13
	"go.miragespace.co/specter/spec/protocol"
14
 
15
	"github.com/caddyserver/certmagic"
16
	"github.com/zhangyunhao116/skipmap"
17
	"go.uber.org/zap"
18
)
19
 
20
type ChordStorage struct {
21
	Logger *zap.Logger
22
	KV     chord.KV
23
 
24
	leaseToken      *skipmap.StringMap[*leaseHolder]
25
	retryInterval   time.Duration
26
	leaseTTL        time.Duration
27
	renewalInterval time.Duration
28
	pollInterval    time.Duration
29
}
30
 
31
type StorageConfig struct {
32
	RetryInterval time.Duration
33
	LeaseTTL      time.Duration
34
}
35
 
36
func NewChordStorage(logger *zap.Logger, kv chord.KV, cfg StorageConfig) (*ChordStorage, error) {
37
	// TODO: assert sensible interval
38
	return &ChordStorage{
39
		Logger:          logger,
40
		KV:              kv,
41
		leaseToken:      skipmap.NewString[*leaseHolder](),
42
		retryInterval:   cfg.RetryInterval,
43
		leaseTTL:        cfg.LeaseTTL,
44
		renewalInterval: cfg.LeaseTTL / 4,
45
		pollInterval:    cfg.LeaseTTL / 2,
46
	}, nil
47
}
48
 
49
func (c *ChordStorage) Lock(ctx context.Context, key string) error {
50
	// c.Logger.Debug("Lock invoked", zap.String("key", key))
51
	for {
52
		token, err := c.KV.Acquire(ctx, []byte(kvKeyName(key)), c.leaseTTL)
53
		switch err {
54
		case chord.ErrKVLeaseConflict:
55
			c.Logger.Debug("Lease acquire conflict, retrying", zap.String("key", key))
56
			<-time.After(c.pollInterval)
57
			continue
58
		case nil:
59
			return c.startLeaseRenewal(key, token)
60
		default:
61
			c.Logger.Error("Error acquiring lease", zap.String("key", key), zap.Error(err))
62
			return err
63
		}
64
	}
65
}
66
 
67
func (c *ChordStorage) renewLease(key string, l *leaseHolder) {
68
	ticker := time.NewTicker(c.renewalInterval)
69
	defer ticker.Stop()
70
 
71
	defer l.Done()
72
 
73
	for {
74
		select {
75
		case <-l.ctx.Done():
76
			return
77
		case <-ticker.C:
78
			if err := c.renewLeaseOnce(context.Background(), key, l); err != nil {
79
				c.Logger.Error("failed to renew lease", zap.String("lease", key), zap.Error(err))
80
				return
81
			}
82
		}
83
	}
84
}
85
 
86
func (c *ChordStorage) startLeaseRenewal(key string, token uint64) error {
87
	leaseCtx, leaseCancel := context.WithCancel(context.Background())
88
	h := &leaseHolder{
89
		token:    token,
90
		ctx:      leaseCtx,
91
		cancelFn: leaseCancel,
92
	}
93
	h.Add(1)
94
	c.leaseToken.Store(key, h)
95
	go c.renewLease(key, h)
96
 
97
	c.Logger.Debug("Lease acquired", zap.String("key", key), zap.Uint64("lease", token))
98
	return nil
99
}
100
 
101
func (c *ChordStorage) renewLeaseOnce(ctx context.Context, key string, l *leaseHolder) error {
102
	prev := atomic.LoadUint64(&l.token)
103
	next, err := c.KV.Renew(ctx, []byte(kvKeyName(key)), c.leaseTTL, prev)
104
	if err != nil {
105
		return err
106
	}
107
	c.Logger.Debug("Lease renewal", zap.String("key", key), zap.Uint64("newToken", next))
108
	atomic.StoreUint64(&l.token, next)
109
	return nil
110
}
111
 
112
func (c *ChordStorage) RenewLockLease(ctx context.Context, key string, leaseDuration time.Duration) error {
113
	lease, ok := c.leaseToken.Load(key)
114
	if !ok {
115
		return fmt.Errorf("not a lease holder of key %s", key)
116
	}
117
	return c.renewLeaseOnce(ctx, key, lease)
118
}
119
 
120
func (c *ChordStorage) Unlock(ctx context.Context, key string) error {
121
	lease, ok := c.leaseToken.LoadAndDelete(key)
122
	if !ok {
123
		return fmt.Errorf("not a lease holder of key %s", key)
124
	}
125
	lease.cancelFn()
126
	lease.Wait()
127
	c.Logger.Debug("Lease released", zap.String("key", key))
128
	return c.KV.Release(ctx, []byte(kvKeyName(key)), atomic.LoadUint64(&lease.token))
129
}
130
 
131
func (c *ChordStorage) Store(ctx context.Context, key string, value []byte) error {
132
	// c.Logger.Debug("Store invoked", zap.String("key", key))
133
	return c.KV.Put(ctx, []byte(kvKeyName(key)), value)
134
}
135
 
136
func (c *ChordStorage) Load(ctx context.Context, key string) ([]byte, error) {
137
	val, err := c.KV.Get(ctx, []byte(kvKeyName(key)))
138
	if err != nil {
139
		return nil, err
140
	}
141
	if val == nil {
142
		c.Logger.Debug("Load returned not found", zap.String("key", key))
143
		return nil, fs.ErrNotExist
144
	}
145
	c.Logger.Debug("Load returned something", zap.String("key", key), zap.Int("val_length", len(val)))
146
	return val, nil
147
}
148
 
149
func (c *ChordStorage) Delete(ctx context.Context, key string) error {
150
	// c.Logger.Debug("Delete invoked", zap.String("key", key))
151
	return c.KV.Delete(ctx, []byte(kvKeyName(key)))
152
}
153
 
154
func (c *ChordStorage) Exists(ctx context.Context, key string) bool {
155
	// c.Logger.Debug("Exists invoked", zap.String("key", key))
156
	val, err := c.KV.Get(ctx, []byte(kvKeyName(key)))
157
	if err != nil {
158
		c.Logger.Debug("Exists error", zap.String("key", key), zap.Error(err))
159
		return false
160
	}
161
	c.Logger.Debug("Exists", zap.String("key", key), zap.Bool("exists", val != nil))
162
	return val != nil
163
}
164
 
165
func (c *ChordStorage) List(ctx context.Context, prefix string, recursive bool) ([]string, error) {
166
	oldPrefix := prefix
167
 
168
	prefix = kvKeyName(prefix)
169
 
170
	keys, err := c.KV.ListKeys(ctx, []byte(prefix))
171
	if err != nil {
172
		return nil, err
173
	}
174
 
175
	var newKey string
176
	found := make([]string, 0)
177
	if recursive {
178
		for _, key := range keys {
179
			if key.GetType() != protocol.KeyComposite_SIMPLE {
180
				continue
181
			}
182
			newKey = strings.TrimPrefix(string(key.GetKey()), kvKeyPrefix)
183
			found = append(found, newKey)
184
		}
185
	} else {
186
		seen := make(map[string]bool)
187
		if !strings.HasSuffix(prefix, "/") {
188
			prefix += "/"
189
		}
190
		for _, key := range keys {
191
			if key.GetType() != protocol.KeyComposite_SIMPLE {
192
				continue
193
			}
194
			sub := strings.TrimPrefix(string(key.GetKey()), prefix)
195
			before, _, ok := strings.Cut(sub, "/")
196
 
197
			if !ok {
198
				newKey = string(key.GetKey())
199
			} else {
200
				newKey = prefix + before
201
			}
202
			newKey = strings.TrimPrefix(newKey, kvKeyPrefix)
203
 
204
			if ok := seen[newKey]; ok {
205
				continue
206
			}
207
			seen[newKey] = true
208
			found = append(found, newKey)
209
		}
210
	}
211
 
212
	c.Logger.Debug("List", zap.String("prefix", oldPrefix), zap.Strings("keys", found), zap.Bool("recursive", recursive))
213
 
214
	return found, nil
215
}
216
 
217
func (c *ChordStorage) Stat(ctx context.Context, key string) (certmagic.KeyInfo, error) {
218
	// c.Logger.Debug("Stat invoked", zap.String("key", key))
219
	info := certmagic.KeyInfo{}
220
	value, err := c.KV.Get(ctx, []byte(kvKeyName(key)))
221
	if err != nil {
222
		return info, err
223
	}
224
	if value == nil {
225
		return info, fs.ErrNotExist
226
	}
227
	info.IsTerminal = true
228
	info.Size = int64(len(value))
229
	info.Key = key
230
	return info, nil
231
}
232
 
233
type leaseHolder struct {
234
	sync.WaitGroup
235
	ctx      context.Context
236
	cancelFn context.CancelFunc
237
	token    uint64
238
}
239
 
240
var _ certmagic.Storage = (*ChordStorage)(nil)
241
var _ certmagic.LockLeaseRenewer = (*ChordStorage)(nil)