| 1 | package dbreads |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "crypto/ecdh" |
| 7 | "crypto/sha256" |
| 8 | "encoding/hex" |
| 9 | "errors" |
| 10 | "fmt" |
| 11 | "hash" |
| 12 | "sync" |
| 13 | "time" |
| 14 | |
| 15 | "github.com/carlosframework/platform/internal/release" |
| 16 | ) |
| 17 | |
| 18 | var errCapped = errors.New("dbreads: the output cap was reached") |
| 19 | |
| 20 | // putCreateOnly creates key with body, retrying until it lands or ctx ends. |
| 21 | // An error or a conflict is resolved by reading the key back: the same |
| 22 | // bytes mean an earlier attempt landed (spec §5.4). The caller keeps body |
| 23 | // until this returns, so a retry is the same bytes. |
| 24 | func putCreateOnly(ctx context.Context, s release.Store, key string, body []byte) error { |
| 25 | wait := 250 * time.Millisecond |
| 26 | for { |
| 27 | _, err := s.Put(ctx, key, body, "") |
| 28 | if err == nil { |
| 29 | return nil |
| 30 | } |
| 31 | got, _, gerr := s.Get(ctx, key) |
| 32 | switch { |
| 33 | case gerr == nil && bytes.Equal(got, body): |
| 34 | return nil |
| 35 | case gerr == nil: |
| 36 | return fmt.Errorf("dbreads: %s holds other bytes: %w", key, release.ErrCASConflict) |
| 37 | } |
| 38 | select { |
| 39 | case <-ctx.Done(): |
| 40 | return ctx.Err() |
| 41 | case <-time.After(wait): |
| 42 | } |
| 43 | wait = min(wait*2, 5*time.Second) |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | // chunkWriter seals a read's output into out/<seq> objects: one every |
| 48 | // size bytes (ChunkBytes for a query, PullChunkBytes for a pull), and whatever is pending every `every`; close seals what is |
| 49 | // left, possibly nothing, as the last chunk. Past max nothing more is kept |
| 50 | // and Write answers errCapped. Each put is create-only and confirmed before |
| 51 | // the next is sealed, under the lock, so seqs are in order (the run |
| 52 | // chunker's discipline). |
| 53 | type chunkWriter struct { |
| 54 | ctx context.Context |
| 55 | store release.Store |
| 56 | id string |
| 57 | pub *ecdh.PublicKey |
| 58 | stream string |
| 59 | max int64 |
| 60 | size int |
| 61 | |
| 62 | mu sync.Mutex |
| 63 | pending []byte |
| 64 | seq int64 |
| 65 | total int64 |
| 66 | sum hash.Hash |
| 67 | err error |
| 68 | isCapped bool |
| 69 | |
| 70 | stop chan struct{} |
| 71 | done chan struct{} |
| 72 | } |
| 73 | |
| 74 | func newChunkWriter(ctx context.Context, s release.Store, id string, pub *ecdh.PublicKey, stream string, every time.Duration, max int64, size int) *chunkWriter { |
| 75 | c := &chunkWriter{ctx: ctx, store: s, id: id, pub: pub, stream: stream, max: max, size: size, sum: sha256.New(), |
| 76 | stop: make(chan struct{}), done: make(chan struct{})} |
| 77 | go c.tick(every) |
| 78 | return c |
| 79 | } |
| 80 | |
| 81 | func (c *chunkWriter) tick(every time.Duration) { |
| 82 | defer close(c.done) |
| 83 | t := time.NewTicker(every) |
| 84 | defer t.Stop() |
| 85 | for { |
| 86 | select { |
| 87 | case <-c.stop: |
| 88 | return |
| 89 | case <-t.C: |
| 90 | c.mu.Lock() |
| 91 | if len(c.pending) > 0 && c.err == nil { |
| 92 | c.err = c.flushLocked(c.pending, false) |
| 93 | c.pending = nil |
| 94 | } |
| 95 | c.mu.Unlock() |
| 96 | } |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | func (c *chunkWriter) Write(p []byte) (int, error) { |
| 101 | c.mu.Lock() |
| 102 | defer c.mu.Unlock() |
| 103 | if c.err != nil { |
| 104 | return 0, c.err |
| 105 | } |
| 106 | if c.isCapped { |
| 107 | return 0, errCapped |
| 108 | } |
| 109 | capped := false |
| 110 | if room := c.max - c.total; int64(len(p)) > room { |
| 111 | p, capped = p[:room], true |
| 112 | } |
| 113 | c.sum.Write(p) |
| 114 | c.total += int64(len(p)) |
| 115 | c.pending = append(c.pending, p...) |
| 116 | for len(c.pending) >= c.size { |
| 117 | if err := c.flushLocked(c.pending[:c.size], false); err != nil { |
| 118 | c.err = err |
| 119 | return 0, err |
| 120 | } |
| 121 | c.pending = append([]byte(nil), c.pending[c.size:]...) |
| 122 | } |
| 123 | if capped { |
| 124 | c.isCapped = true |
| 125 | return len(p), errCapped |
| 126 | } |
| 127 | return len(p), nil |
| 128 | } |
| 129 | |
| 130 | // flushLocked seals and creates the next chunk. seq moves only once the |
| 131 | // create is confirmed, so close never reports, and no result ever names, a |
| 132 | // chunk that was not confirmed (spec §5.4; plan review finding 9). A chunk |
| 133 | // whose create stayed ambiguous until the context ended may land later at |
| 134 | // seq+1; no result names it, so no reader relays it. |
| 135 | func (c *chunkWriter) flushLocked(b []byte, last bool) error { |
| 136 | next := c.seq + 1 |
| 137 | sealed, err := SealChunk(c.pub, c.id, next, c.stream, last, b) |
| 138 | if err != nil { |
| 139 | return err |
| 140 | } |
| 141 | if err := putCreateOnly(c.ctx, c.store, ChunkKey(c.id, next), sealed); err != nil { |
| 142 | return err |
| 143 | } |
| 144 | c.seq = next |
| 145 | return nil |
| 146 | } |
| 147 | |
| 148 | func (c *chunkWriter) capped() bool { |
| 149 | c.mu.Lock() |
| 150 | defer c.mu.Unlock() |
| 151 | return c.isCapped |
| 152 | } |
| 153 | |
| 154 | // close seals the pending bytes as the last chunk and reports the last seq, |
| 155 | // the plaintext byte count and its SHA-256. |
| 156 | func (c *chunkWriter) close() (int64, int64, string, error) { |
| 157 | close(c.stop) |
| 158 | <-c.done |
| 159 | c.mu.Lock() |
| 160 | defer c.mu.Unlock() |
| 161 | if c.err == nil { |
| 162 | c.err = c.flushLocked(c.pending, true) |
| 163 | c.pending = nil |
| 164 | } |
| 165 | return c.seq, c.total, hex.EncodeToString(c.sum.Sum(nil)), c.err |
| 166 | } |
| 167 | |
amadan