| 1 | package dbreads |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "crypto/ecdh" |
| 7 | "crypto/ed25519" |
| 8 | "crypto/rand" |
| 9 | "crypto/sha256" |
| 10 | "encoding/hex" |
| 11 | "encoding/json" |
| 12 | "errors" |
| 13 | "net" |
| 14 | "os" |
| 15 | "sort" |
| 16 | "sync" |
| 17 | "sync/atomic" |
| 18 | "testing" |
| 19 | "time" |
| 20 | |
| 21 | lsfile "github.com/benbjohnson/litestream/file" |
| 22 | |
| 23 | "github.com/carlosframework/platform/internal/dbdoor" |
| 24 | "github.com/carlosframework/platform/internal/dbdoor/dbdoortest" |
| 25 | "github.com/carlosframework/platform/internal/hostclaim" |
| 26 | "github.com/carlosframework/platform/internal/provision" |
| 27 | "github.com/carlosframework/platform/internal/release" |
| 28 | ) |
| 29 | |
| 30 | const fxAccount, fxApp, fxHost = "bab", "shop", "shop.example.com" |
| 31 | |
| 32 | // inProcWorkers runs the real ServeWorker in process over net.Pipe: not |
| 33 | // acceptance (Task 15 drives the installed units). units counts the live |
| 34 | // workers, the way the box counts carlos-dbread@ units. |
| 35 | type inProcWorkers struct { |
| 36 | t *testing.T |
| 37 | opt WorkerOptions |
| 38 | live atomic.Int64 |
| 39 | dials atomic.Int64 |
| 40 | // linger keeps a finished worker counted as a live unit for a while, |
| 41 | // as systemd's deactivation does; liveAtDial records the live count |
| 42 | // each dial saw. |
| 43 | linger time.Duration |
| 44 | mu sync.Mutex |
| 45 | liveAtDial []int64 |
| 46 | // failStart makes the worker end at once without a word, as a unit |
| 47 | // whose service fails to start after the socket accepted does. |
| 48 | failStart atomic.Bool |
| 49 | // sent is the type of every frame the agent wrote, in order. |
| 50 | sent []byte |
| 51 | } |
| 52 | |
| 53 | // frameTap records the first byte of each Write the agent makes: one frame |
| 54 | // type per frame (writeFrame writes a frame in one Write). |
| 55 | type frameTap struct { |
| 56 | net.Conn |
| 57 | w *inProcWorkers |
| 58 | } |
| 59 | |
| 60 | func (f frameTap) Write(p []byte) (int, error) { |
| 61 | if len(p) > 0 { |
| 62 | f.w.mu.Lock() |
| 63 | f.w.sent = append(f.w.sent, p[0]) |
| 64 | f.w.mu.Unlock() |
| 65 | } |
| 66 | return f.Conn.Write(p) |
| 67 | } |
| 68 | |
| 69 | func (w *inProcWorkers) sentFrames() []byte { |
| 70 | w.mu.Lock() |
| 71 | defer w.mu.Unlock() |
| 72 | return append([]byte(nil), w.sent...) |
| 73 | } |
| 74 | |
| 75 | func (w *inProcWorkers) dial(context.Context) (net.Conn, error) { |
| 76 | agentEnd, workerEnd := net.Pipe() |
| 77 | w.mu.Lock() |
| 78 | w.liveAtDial = append(w.liveAtDial, w.live.Load()) |
| 79 | w.mu.Unlock() |
| 80 | w.dials.Add(1) |
| 81 | w.live.Add(1) |
| 82 | if w.failStart.Load() { |
| 83 | // The socket accepted and the service died: the plan's header is |
| 84 | // read by nobody useful, and the connection drops. |
| 85 | go func() { |
| 86 | readFrame(workerEnd) |
| 87 | workerEnd.Close() |
| 88 | time.Sleep(w.linger) |
| 89 | w.live.Add(-1) |
| 90 | }() |
| 91 | return frameTap{Conn: agentEnd, w: w}, nil |
| 92 | } |
| 93 | opt := w.opt |
| 94 | if opt.Scratch == "" { |
| 95 | opt.Scratch = w.t.TempDir() |
| 96 | } |
| 97 | if opt.Exit == nil { |
| 98 | opt.Exit = func(int) {} |
| 99 | } |
| 100 | if opt.Self == nil { |
| 101 | opt.Self = func() HelloReply { return HelloReply{V: ProtocolVersion, Interfaces: []string{"lo"}} } |
| 102 | } |
| 103 | go func() { |
| 104 | ServeWorker(workerEnd, opt) |
| 105 | workerEnd.Close() |
| 106 | time.Sleep(w.linger) |
| 107 | w.live.Add(-1) |
| 108 | }() |
| 109 | return frameTap{Conn: agentEnd, w: w}, nil |
| 110 | } |
| 111 | |
| 112 | func (w *inProcWorkers) units(context.Context) (int, error) { return int(w.live.Load()), nil } |
| 113 | |
| 114 | func plentyFree() (uint64, uint64, error) { return 100 << 30, 200 << 30, nil } |
| 115 | |
| 116 | type fixture struct { |
| 117 | t *testing.T |
| 118 | store release.Store |
| 119 | rootPub ed25519.PublicKey |
| 120 | rootPriv ed25519.PrivateKey |
| 121 | reply *ecdh.PrivateKey |
| 122 | replica string |
| 123 | now time.Time |
| 124 | } |
| 125 | |
| 126 | // newFixture is a deployment store holding shop.example.com's record in uk, |
| 127 | // and a file replica of its database with batches*10 rows. |
| 128 | func newFixture(t *testing.T, batches int) *fixture { |
| 129 | t.Helper() |
| 130 | pub, priv, err := ed25519.GenerateKey(rand.Reader) |
| 131 | if err != nil { |
| 132 | t.Fatal(err) |
| 133 | } |
| 134 | reply, err := NewReplyKey() |
| 135 | if err != nil { |
| 136 | t.Fatal(err) |
| 137 | } |
| 138 | fx := &fixture{t: t, store: release.NewDirStore(t.TempDir()), rootPub: pub, rootPriv: priv, reply: reply, |
| 139 | replica: t.TempDir(), now: time.Now().UTC().Truncate(time.Second)} |
| 140 | if batches > 0 { |
| 141 | dbdoortest.Replicate(t, fx.replica, batches) |
| 142 | } |
| 143 | fx.putRecord(provision.Record{Host: fxHost, Region: "uk", CreatedAt: time.Date(2026, 9, 1, 0, 0, 0, 0, time.UTC)}) |
| 144 | return fx |
| 145 | } |
| 146 | |
| 147 | func (fx *fixture) putRecordUnder(acct string, rec provision.Record) { |
| 148 | fx.t.Helper() |
| 149 | rec.Account, rec.App = acct, fxApp |
| 150 | if rec.Host == "" { |
| 151 | rec.Host = fxHost |
| 152 | } |
| 153 | if rec.Hibernate == nil { |
| 154 | yes := true |
| 155 | rec.Hibernate = &yes |
| 156 | } |
| 157 | body, err := json.Marshal(rec) |
| 158 | if err != nil { |
| 159 | fx.t.Fatal(err) |
| 160 | } |
| 161 | key := provision.RecordKey(acct, fxApp, rec.Host) |
| 162 | if err := fx.store.Delete(context.Background(), key); err != nil { |
| 163 | fx.t.Fatal(err) |
| 164 | } |
| 165 | if _, err := fx.store.Put(context.Background(), key, body, ""); err != nil { |
| 166 | fx.t.Fatal(err) |
| 167 | } |
| 168 | } |
| 169 | |
| 170 | func (fx *fixture) putRecord(rec provision.Record) { fx.putRecordUnder(fxAccount, rec) } |
| 171 | |
| 172 | func (fx *fixture) record() provision.Record { |
| 173 | fx.t.Helper() |
| 174 | rec, err := provision.GetRecord(context.Background(), fx.store, fxAccount, fxApp, fxHost) |
| 175 | if err != nil { |
| 176 | fx.t.Fatal(err) |
| 177 | } |
| 178 | return *rec |
| 179 | } |
| 180 | |
| 181 | // request is the envelope the console would sign for the current record. |
| 182 | func (fx *fixture) request(sqlText, at string) Request { |
| 183 | fx.t.Helper() |
| 184 | rec := fx.record() |
| 185 | inc, err := hostclaim.Incarnation(rec.Account, rec.App, rec.Host, rec.CreatedAt) |
| 186 | if err != nil { |
| 187 | fx.t.Fatal(err) |
| 188 | } |
| 189 | id, err := MintID() |
| 190 | if err != nil { |
| 191 | fx.t.Fatal(err) |
| 192 | } |
| 193 | return Request{ID: id, Account: fxAccount, App: fxApp, Host: fxHost, Incarnation: inc, Region: "uk", |
| 194 | Op: OpQuery, SQL: sqlText, At: at, ReplyKey: EncodeReplyKey(fx.reply.PublicKey()), |
| 195 | By: "member@example.com", IssuedAt: fx.now.Format(time.RFC3339)} |
| 196 | } |
| 197 | |
| 198 | // submit is the console's two creates: the signed request, then the queue |
| 199 | // entry. |
| 200 | func (fx *fixture) submit(req Request) QueueEntry { |
| 201 | fx.t.Helper() |
| 202 | wire, err := Sign(release.RootSigner(fx.rootPriv), req) |
| 203 | if err != nil { |
| 204 | fx.t.Fatal(err) |
| 205 | } |
| 206 | if _, err := fx.store.Put(context.Background(), RequestKey(req.ID), wire, ""); err != nil { |
| 207 | fx.t.Fatal(err) |
| 208 | } |
| 209 | issued, _ := req.Issued() |
| 210 | key := QueueKey(req.Region, issued, req.ID) |
| 211 | if _, err := fx.store.Put(context.Background(), key, nil, ""); err != nil { |
| 212 | fx.t.Fatal(err) |
| 213 | } |
| 214 | e, _ := ParseQueueKey(key) |
| 215 | return e |
| 216 | } |
| 217 | |
| 218 | // agent is a capable agent in uk over the fixture's store and replica. |
| 219 | // Its probe is Task 12's: it starts capable. |
| 220 | func (fx *fixture) agent(box string, mut ...func(*AgentConfig, *inProcWorkers)) (*Agent, *inProcWorkers) { |
| 221 | w := &inProcWorkers{t: fx.t} |
| 222 | cfg := AgentConfig{ |
| 223 | Store: fx.store, Region: "uk", Box: box, Boot: "boot-" + box, Pub: fx.rootPub, HibernatePrefix: "instances", |
| 224 | OpenReplica: func(_ context.Context, prefix string) (dbdoor.ReplicaReader, error) { |
| 225 | if prefix != "instances/"+fxHost { |
| 226 | fx.t.Errorf("the agent asked for prefix %q", prefix) |
| 227 | } |
| 228 | return lsfile.NewReplicaClient(fx.replica), nil |
| 229 | }, |
| 230 | Units: w.units, Free: plentyFree, Logf: fx.t.Logf, |
| 231 | Jitter: func() time.Duration { return 0 }, |
| 232 | } |
| 233 | for _, m := range mut { |
| 234 | m(&cfg, w) |
| 235 | } |
| 236 | cfg.DialWorker = w.dial |
| 237 | a := NewAgent(cfg) |
| 238 | a.probeOK.Store(true) |
| 239 | return a, w |
| 240 | } |
| 241 | |
| 242 | // chunkKeys lists a read's chunk objects. |
| 243 | func (fx *fixture) chunkKeys(id string) []string { |
| 244 | keys, err := fx.store.List(context.Background(), OutPrefix(id)) |
| 245 | if err != nil && !errors.Is(err, os.ErrNotExist) { |
| 246 | fx.t.Fatal(err) |
| 247 | } |
| 248 | sort.Strings(keys) |
| 249 | return keys |
| 250 | } |
| 251 | |
| 252 | // answer reads a read's result and opens its chunks as the CLI would. |
| 253 | func (fx *fixture) answer(id string) ([]byte, Summary, Result) { |
| 254 | fx.t.Helper() |
| 255 | raw, _, err := fx.store.Get(context.Background(), ResultKey(id)) |
| 256 | if err != nil { |
| 257 | fx.t.Fatalf("no result for %s: %v", id, err) |
| 258 | } |
| 259 | var r Result |
| 260 | if err := json.Unmarshal(raw, &r); err != nil || r.Valid() != nil { |
| 261 | fx.t.Fatalf("result %s: %v", raw, err) |
| 262 | } |
| 263 | s, err := OpenSummary(fx.reply, id, r.Sealed) |
| 264 | if err != nil { |
| 265 | fx.t.Fatal(err) |
| 266 | } |
| 267 | var out bytes.Buffer |
| 268 | for seq := int64(1); seq <= r.LastSeq; seq++ { |
| 269 | sealed, _, err := fx.store.Get(context.Background(), ChunkKey(id, seq)) |
| 270 | if err != nil { |
| 271 | fx.t.Fatalf("chunk %d: %v", seq, err) |
| 272 | } |
| 273 | plain, err := OpenChunk(fx.reply, id, seq, StreamRows, seq == r.LastSeq, sealed) |
| 274 | if err != nil { |
| 275 | fx.t.Fatalf("chunk %d: %v", seq, err) |
| 276 | } |
| 277 | out.Write(plain) |
| 278 | } |
| 279 | sum := sha256.Sum256(out.Bytes()) |
| 280 | if r.State == StateDone && (int64(out.Len()) != s.Bytes || hex.EncodeToString(sum[:]) != s.SHA256) { |
| 281 | fx.t.Fatalf("the answer does not add up: %d bytes, summary %+v", out.Len(), s) |
| 282 | } |
| 283 | return out.Bytes(), s, r |
| 284 | } |
| 285 | |
amadan