From 04fc7d53dd3da6f76e66fcf6e28e8d0c52daff46 Mon Sep 17 00:00:00 2001 From: Artem Lytkin Date: Sun, 18 Oct 2026 18:15:03 +0300 Subject: [PATCH] XDRIVE transport: Carry a Google Drive session in a ring of slots Co-authored-by: Ahmad-2213 <106820058+Ahmad-2213@users.noreply.github.com> --- transport/internet/xdrive/conn.go | 28 +- transport/internet/xdrive/drive.go | 152 ++-- transport/internet/xdrive/drive_live_test.go | 10 +- transport/internet/xdrive/drive_test.go | 43 +- transport/internet/xdrive/slots.go | 758 +++++++++++++++++++ transport/internet/xdrive/slots_test.go | 307 ++++++++ transport/internet/xdrive/wal.go | 4 + transport/internet/xdrive/xdrive.go | 53 +- 8 files changed, 1284 insertions(+), 71 deletions(-) create mode 100644 transport/internet/xdrive/slots.go create mode 100644 transport/internet/xdrive/slots_test.go diff --git a/transport/internet/xdrive/conn.go b/transport/internet/xdrive/conn.go index 50f871718759..516db929aad3 100644 --- a/transport/internet/xdrive/conn.go +++ b/transport/internet/xdrive/conn.go @@ -2,6 +2,7 @@ package xdrive import ( "context" + "io" "os" "sync" "time" @@ -11,10 +12,12 @@ import ( var placeholderAddr = &net.TCPAddr{IP: net.IP{127, 0, 0, 1}, Port: 0} +// Conn is a session in segments, walWriter and walReader, or in slots, +// slotWriter and slotReader type Conn struct { cancel context.CancelFunc - writer *walWriter - reader *walReader + writer io.WriteCloser + reader chunkReader onClose func() readBuf []byte @@ -27,6 +30,13 @@ type Conn struct { closeErr error } +type chunkReader interface { + // chunks is closed when there are no more, then Err says why + chunks() <-chan []byte + Err() error + Wake() +} + func newConn(ctx context.Context, storage Storage, writePrefix, readPrefix string, p params, onClose func()) *Conn { ctx, cancel := context.WithCancel(ctx) return &Conn{ @@ -37,6 +47,16 @@ func newConn(ctx context.Context, storage Storage, writePrefix, readPrefix strin } } +func newSlotConn(ctx context.Context, store slotStore, write, read slotTable, p params, onClose func()) *Conn { + ctx, cancel := context.WithCancel(ctx) + return &Conn{ + cancel: cancel, + writer: newSlotWriter(ctx, store, write, p), + reader: newSlotReader(ctx, store, read, p), + onClose: onClose, + } +} + func (c *Conn) Read(b []byte) (int, error) { if len(c.readBuf) == 0 { data, err := c.receive() @@ -53,7 +73,7 @@ func (c *Conn) Read(b []byte) (int, error) { func (c *Conn) receive() ([]byte, error) { deadline := c.getDeadline(true) if deadline.IsZero() { - data, ok := <-c.reader.ch + data, ok := <-c.reader.chunks() if !ok { return nil, c.reader.Err() } @@ -67,7 +87,7 @@ func (c *Conn) receive() ([]byte, error) { defer timer.Stop() select { - case data, ok := <-c.reader.ch: + case data, ok := <-c.reader.chunks(): if !ok { return nil, c.reader.Err() } diff --git a/transport/internet/xdrive/drive.go b/transport/internet/xdrive/drive.go index 472efc91b95c..658b14a5a91a 100644 --- a/transport/internet/xdrive/drive.go +++ b/transport/internet/xdrive/drive.go @@ -401,24 +401,30 @@ func (s *driveStorage) Put(ctx context.Context, name string, data []byte) error } func (s *driveStorage) putInline(ctx context.Context, name string, data []byte) error { + _, err := s.create(ctx, name, base64.StdEncoding.EncodeToString(data)) + return err +} + +// create makes an object with description and no body and returns its id +func (s *driveStorage) create(ctx context.Context, name, description string) (string, error) { flat := flatten(name) body, err := json.Marshal(map[string]interface{}{ "name": flat, "parents": []string{s.folder}, - "description": base64.StdEncoding.EncodeToString(data), + "description": description, }) if err != nil { - return errors.New("failed to build the inline metadata").Base(err) + return "", errors.New("failed to build the inline metadata").Base(err) } status, payload, err := s.do(ctx, http.MethodPost, s.filesURL+"?fields=id", "application/json; charset=UTF-8", body) if err != nil { - return err + return "", err } if status != http.StatusOK { - return errors.New("Drive rejected the inline upload of ", name, + return "", errors.New("Drive rejected the inline upload of ", name, " with ", status, ": ", string(payload)) } @@ -426,12 +432,23 @@ func (s *driveStorage) putInline(ctx context.Context, name string, data []byte) ID string `json:"id"` } if err := json.Unmarshal(payload, &parsed); err != nil { - return errors.New("failed to parse the Drive upload response").Base(err) + return "", errors.New("failed to parse the Drive upload response").Base(err) } if parsed.ID != "" { s.rememberID(flat, parsed.ID) } - return nil + return parsed.ID, nil +} + +// multipartBody is a body of an upload with metadata and data +func multipartBody(metadata, data []byte) []byte { + var body bytes.Buffer + fmt.Fprintf(&body, "--%s\r\nContent-Type: application/json; charset=UTF-8\r\n\r\n", driveBoundary) + body.Write(metadata) + fmt.Fprintf(&body, "\r\n--%s\r\nContent-Type: application/octet-stream\r\n\r\n", driveBoundary) + body.Write(data) + fmt.Fprintf(&body, "\r\n--%s--\r\n", driveBoundary) + return body.Bytes() } func (s *driveStorage) putMedia(ctx context.Context, name string, data []byte) error { @@ -445,15 +462,8 @@ func (s *driveStorage) putMedia(ctx context.Context, name string, data []byte) e return errors.New("failed to build the upload metadata").Base(err) } - var body bytes.Buffer - fmt.Fprintf(&body, "--%s\r\nContent-Type: application/json; charset=UTF-8\r\n\r\n", driveBoundary) - body.Write(metadata) - fmt.Fprintf(&body, "\r\n--%s\r\nContent-Type: application/octet-stream\r\n\r\n", driveBoundary) - body.Write(data) - fmt.Fprintf(&body, "\r\n--%s--\r\n", driveBoundary) - status, payload, err := s.do(ctx, http.MethodPost, s.uploadURL, - "multipart/related; boundary="+driveBoundary, body.Bytes()) + "multipart/related; boundary="+driveBoundary, multipartBody(metadata, data)) if err != nil { return err } @@ -480,45 +490,22 @@ func (s *driveStorage) Get(ctx context.Context, name string) ([]byte, error) { return nil, err } - status, payload, err := s.do(ctx, http.MethodGet, - s.filesURL+"/"+url.PathEscape(id)+"?alt=media&supportsAllDrives=true", "", nil) - if err != nil { - return nil, err - } - switch status { - case http.StatusOK: - if len(payload) > 0 { - return payload, nil - } - return s.getInline(ctx, id) - case http.StatusNotFound: + payload, err := s.readBody(ctx, id) + if err == errNotFound { s.forgetID(flat) - return nil, errNotFound - default: - return nil, errors.New("Drive download of ", name, " answered ", status, ": ", string(payload)) } + if err != nil || len(payload) > 0 { + return payload, err + } + return s.getInline(ctx, id) } func (s *driveStorage) getInline(ctx context.Context, id string) ([]byte, error) { - status, payload, err := s.do(ctx, http.MethodGet, - s.filesURL+"/"+url.PathEscape(id)+"?fields=description&supportsAllDrives=true", "", nil) - if err != nil { + description, err := s.readDescription(ctx, id) + if err != nil || description == "" { return nil, err } - if status != http.StatusOK { - return nil, errors.New("Drive answered ", status, " for inline data: ", string(payload)) - } - - var parsed struct { - Description string `json:"description"` - } - if err := json.Unmarshal(payload, &parsed); err != nil { - return nil, errors.New("failed to parse the inline data").Base(err) - } - if parsed.Description == "" { - return nil, nil - } - data, err := base64.StdEncoding.DecodeString(parsed.Description) + data, err := base64.StdEncoding.DecodeString(description) if err != nil { return nil, errors.New("the inline data is not valid base64").Base(err) } @@ -605,3 +592,76 @@ func (s *driveStorage) Close() error { } return nil } + +// makeSlot implements slotStore +func (s *driveStorage) makeSlot(ctx context.Context, name string) (string, error) { + id, err := s.create(ctx, name, "") + if err == nil && id == "" { + err = errors.New("Drive returned no id for ", name) + } + return id, err +} + +// setSlot implements slotStore, a description and a body go in one request +func (s *driveStorage) setSlot(ctx context.Context, id, description string, body []byte) error { + metadata, err := json.Marshal(map[string]string{"description": description}) + if err != nil { + return errors.New("failed to build the slot metadata").Base(err) + } + target := s.filesURL + "/" + url.PathEscape(id) + "?fields=id&supportsAllDrives=true" + contentType := "application/json; charset=UTF-8" + if body != nil { + upload, _, _ := strings.Cut(s.uploadURL, "?") + target = upload + "/" + url.PathEscape(id) + "?uploadType=multipart&fields=id&supportsAllDrives=true" + contentType = "multipart/related; boundary=" + driveBoundary + metadata = multipartBody(metadata, body) + } + + status, payload, err := s.do(ctx, http.MethodPatch, target, contentType, metadata) + switch { + case err != nil: + return err + case status == http.StatusNotFound: + return errNotFound + case status != http.StatusOK: + return errors.New("Drive answered ", status, " for the update of ", id, ": ", string(payload)) + } + return nil +} + +// readDescription implements slotStore +func (s *driveStorage) readDescription(ctx context.Context, id string) (string, error) { + status, payload, err := s.do(ctx, http.MethodGet, + s.filesURL+"/"+url.PathEscape(id)+"?fields=description&supportsAllDrives=true", "", nil) + switch { + case err != nil: + return "", err + case status == http.StatusNotFound: + return "", errNotFound + case status != http.StatusOK: + return "", errors.New("Drive answered ", status, " for the description of ", id, ": ", string(payload)) + } + + var parsed struct { + Description string `json:"description"` + } + if err := json.Unmarshal(payload, &parsed); err != nil { + return "", errors.New("failed to parse the description of ", id).Base(err) + } + return parsed.Description, nil +} + +// readBody implements slotStore +func (s *driveStorage) readBody(ctx context.Context, id string) ([]byte, error) { + status, payload, err := s.do(ctx, http.MethodGet, + s.filesURL+"/"+url.PathEscape(id)+"?alt=media&supportsAllDrives=true", "", nil) + switch { + case err != nil: + return nil, err + case status == http.StatusNotFound: + return nil, errNotFound + case status != http.StatusOK: + return nil, errors.New("Drive answered ", status, " for the body of ", id, ": ", string(payload)) + } + return payload, nil +} diff --git a/transport/internet/xdrive/drive_live_test.go b/transport/internet/xdrive/drive_live_test.go index f652cd794e0a..7808f8093a9f 100644 --- a/transport/internet/xdrive/drive_live_test.go +++ b/transport/internet/xdrive/drive_live_test.go @@ -57,7 +57,7 @@ func liveDriveConfig(t *testing.T) *Config { config.Concurrency = uint32(envInt(t, "XRAY_XDRIVE_LIVE_CONCURRENCY")) } - storage, err := newDriveStorage(nil, config) + storage, err := newDriveStorage(&internet.MemoryStreamConfig{}, config) if err != nil { t.Fatalf("newDriveStorage: %v", err) } @@ -83,7 +83,7 @@ func envInt(t *testing.T, name string) int { func TestLiveDriveStorage(t *testing.T) { config := liveDriveConfig(t) - storage, err := newDriveStorage(nil, config) + storage, err := newDriveStorage(&internet.MemoryStreamConfig{}, config) if err != nil { t.Fatalf("newDriveStorage: %v", err) } @@ -194,7 +194,7 @@ func TestLiveDriveTransport(t *testing.T) { func TestLiveDriveParallelPut(t *testing.T) { config := liveDriveConfig(t) - storage, err := newDriveStorage(nil, config) + storage, err := newDriveStorage(&internet.MemoryStreamConfig{}, config) if err != nil { t.Fatalf("newDriveStorage: %v", err) } @@ -266,7 +266,7 @@ func TestLiveDriveParallelPut(t *testing.T) { func TestLiveDriveSegmentSweep(t *testing.T) { config := liveDriveConfig(t) - storage, err := newDriveStorage(nil, config) + storage, err := newDriveStorage(&internet.MemoryStreamConfig{}, config) if err != nil { t.Fatalf("newDriveStorage: %v", err) } @@ -303,7 +303,7 @@ func TestLiveDriveSegmentSweep(t *testing.T) { func TestLiveDriveListLag(t *testing.T) { config := liveDriveConfig(t) - storage, err := newDriveStorage(nil, config) + storage, err := newDriveStorage(&internet.MemoryStreamConfig{}, config) if err != nil { t.Fatalf("newDriveStorage: %v", err) } diff --git a/transport/internet/xdrive/drive_test.go b/transport/internet/xdrive/drive_test.go index 7a5a5b888c05..a48fb0b9acac 100644 --- a/transport/internet/xdrive/drive_test.go +++ b/transport/internet/xdrive/drive_test.go @@ -63,6 +63,7 @@ func newFakeDrive(t *testing.T) *fakeDrive { } mux.HandleFunc("/token", record(drive.handleToken)) mux.HandleFunc("/upload", record(drive.handleUpload)) + mux.HandleFunc("/upload/", record(drive.handleUpload)) mux.HandleFunc("/files", record(drive.handleFiles)) mux.HandleFunc("/files/", record(drive.handleFile)) drive.server = httptest.NewServer(mux) @@ -119,7 +120,8 @@ func (d *fakeDrive) handleUpload(w http.ResponseWriter, r *http.Request) { return } var metadata struct { - Name string `json:"name"` + Name string `json:"name"` + Description string `json:"description"` } if err := json.NewDecoder(metaPart).Decode(&metadata); err != nil { w.WriteHeader(http.StatusBadRequest) @@ -131,6 +133,22 @@ func (d *fakeDrive) handleUpload(w http.ResponseWriter, r *http.Request) { data, _ = io.ReadAll(dataPart) } + if id, ok := strings.CutPrefix(r.URL.Path, "/upload/"); ok { + // an update of the body and description of a file + d.mu.Lock() + file := d.files[id] + if file != nil { + file.data, file.description = data, metadata.Description + } + d.mu.Unlock() + if file == nil || r.Method != http.MethodPatch { + w.WriteHeader(http.StatusNotFound) + return + } + json.NewEncoder(w).Encode(map[string]string{"id": id}) + return + } + d.mu.Lock() d.nextID++ id := fmt.Sprintf("id-%d", d.nextID) @@ -223,6 +241,11 @@ func (d *fakeDrive) handleFile(w http.ResponseWriter, r *http.Request) { if ok && r.Method == http.MethodDelete { delete(d.files, id) } + var snapshot fakeFile + if ok { + // a file can be updated, its fields are read here + snapshot = *file + } d.mu.Unlock() if !ok { @@ -233,11 +256,25 @@ func (d *fakeDrive) handleFile(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNoContent) return } + if r.Method == http.MethodPatch { + var meta struct { + Description string `json:"description"` + } + if err := json.NewDecoder(r.Body).Decode(&meta); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + d.mu.Lock() + file.description = meta.Description + d.mu.Unlock() + json.NewEncoder(w).Encode(map[string]string{"id": id}) + return + } if strings.Contains(r.URL.RawQuery, "fields=description") { - json.NewEncoder(w).Encode(map[string]string{"description": file.description}) + json.NewEncoder(w).Encode(map[string]string{"description": snapshot.description}) return } - w.Write(file.data) + w.Write(snapshot.data) } func (d *fakeDrive) shouldFail(kind string) bool { diff --git a/transport/internet/xdrive/slots.go b/transport/internet/xdrive/slots.go new file mode 100644 index 000000000000..8417acb99d9b --- /dev/null +++ b/transport/internet/xdrive/slots.go @@ -0,0 +1,758 @@ +package xdrive + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/binary" + "encoding/json" + "hash/crc32" + "io" + "strconv" + "sync" + "time" + + "github.com/xtls/xray-core/common/errors" +) + +// A session on a slotStore is carried in a ring of slots for each direction, +// objects that both peers know by id: segment n is written to slot n % depth +// and read from it, with nothing to list and no index to wait for. The reader +// of a direction keeps how far it has read in a watermark, and the writer +// reuses a slot only once that says what was in it before was taken + +const ( + slotMagic = "XDRVSL01" + slotHeaderLen = len(slotMagic) + 8 + 1 + 4 + 4 + // slotInlineLimit is how large a segment is at most to go in the description + // of its slot, which takes one request less than a body + slotInlineLimit = 14 << 10 + slotDepthLimit = 8 + // inlineRampSegments segments of one size in a row double it, from + // slotInlineLimit up to segmentBytes + inlineRampSegments = 8 + + slotInline = 1 + slotMedia = 2 + slotEnd = 3 + slotErr = 4 +) + +var ( + // the head is asked for between slotHuntFloor and slotHuntCeiling, a segment + // behind it at about the pace that segments come + slotHuntFloor = 50 * time.Millisecond + slotHuntCeiling = 200 * time.Millisecond + slotTailMin = 150 * time.Millisecond + slotTailMax = 800 * time.Millisecond + slotFetchAhead = 6 + + slotAckRetry = 100 * time.Millisecond + slotWatermarkTTL = 250 * time.Millisecond + slotStallTimeout = 60 * time.Second + slotCloseGrace = 10 * time.Second + slotGoneTimeout = 60 * time.Second +) + +// slotStore is a storage of objects that are written and read again by id, +// each with a description and a body +type slotStore interface { + makeSlot(ctx context.Context, name string) (string, error) + // setSlot sets the description of a slot and, unless body is nil, its body + setSlot(ctx context.Context, id, description string, body []byte) error + readDescription(ctx context.Context, id string) (string, error) + readBody(ctx context.Context, id string) ([]byte, error) +} + +// slotFrame is what the description of a slot says: which segment it holds, +// and either the segment itself or the size and checksum of the body that does +type slotFrame struct { + seq int64 + mode byte + size int + crc uint32 + data []byte +} + +func newSlotFrame(seq int64, mode byte, data []byte) slotFrame { + f := slotFrame{seq: seq, mode: mode, size: len(data), crc: crc32.ChecksumIEEE(data)} + if mode == slotInline { + f.data = data + } + return f +} + +func (f slotFrame) encode() string { + b := make([]byte, slotHeaderLen, slotHeaderLen+len(f.data)) + copy(b, slotMagic) + binary.BigEndian.PutUint64(b[8:], uint64(f.seq)) + b[16] = f.mode + binary.BigEndian.PutUint32(b[17:], uint32(f.size)) + binary.BigEndian.PutUint32(b[21:], f.crc) + return base64.StdEncoding.EncodeToString(append(b, f.data...)) +} + +func decodeSlot(description string) (slotFrame, bool) { + b, err := base64.StdEncoding.DecodeString(description) + if err != nil || len(b) < slotHeaderLen || string(b[:len(slotMagic)]) != slotMagic { + return slotFrame{}, false + } + f := slotFrame{ + seq: int64(binary.BigEndian.Uint64(b[8:])), + mode: b[16], + size: int(binary.BigEndian.Uint32(b[17:])), + crc: binary.BigEndian.Uint32(b[21:]), + } + if f.mode == slotInline { + f.data = b[slotHeaderLen:] + if !f.holds(f.data) { + return slotFrame{}, false + } + } + return f, true +} + +// holds reports whether body is the one that f is about, which a body read +// right after its description need not be yet +func (f slotFrame) holds(body []byte) bool { + return len(body) == f.size && crc32.ChecksumIEEE(body) == f.crc +} + +// slotTable is the ring of one direction: the ids of its slots and of the +// watermark of its reader +type slotTable struct { + Slots []string `json:"s"` + Watermark string `json:"w"` +} + +// ring is what a client announces a session with, the tables are named from +// its side +type ring struct { + Up slotTable `json:"u"` + Down slotTable `json:"d"` +} + +func decodeRing(b []byte) (ring, bool) { + var r ring + if json.Unmarshal(b, &r) != nil || len(r.Up.Slots) == 0 || len(r.Down.Slots) == 0 || + r.Up.Watermark == "" || r.Down.Watermark == "" { + return ring{}, false + } + return r, true +} + +// makeRing makes the slots and watermarks of a session, all at once +func makeRing(ctx context.Context, store slotStore, session string, p params) (ring, error) { + depth := min(max(p.concurrency, 2), slotDepthLimit) + r := ring{ + Up: slotTable{Slots: make([]string, depth)}, + Down: slotTable{Slots: make([]string, depth)}, + } + var ( + wg sync.WaitGroup + errOnce sync.Once + err error + ) + create := func(id *string, name string) { + wg.Add(1) + go func() { + defer wg.Done() + var e error + if *id, e = store.makeSlot(ctx, name); e != nil { + errOnce.Do(func() { err = e }) + } + }() + } + for _, t := range []struct { + table *slotTable + prefix string + }{{&r.Up, uplinkPrefix(session)}, {&r.Down, downlinkPrefix(session)}} { + for i := range t.table.Slots { + create(&t.table.Slots[i], t.prefix+"/s"+strconv.Itoa(i)) + } + create(&t.table.Watermark, t.prefix+"/w") + } + wg.Wait() + return r, err +} + +func readWatermark(ctx context.Context, store slotStore, id string) (int64, error) { + d, err := store.readDescription(ctx, id) + if err != nil || d == "" { + return -1, err + } + return strconv.ParseInt(d, 10, 64) +} + +type slotWriter struct { + ctx context.Context + store slotStore + table slotTable + params + // slots has a token for each segment on its way, at most one for each slot + slots chan struct{} + wg sync.WaitGroup + + mu sync.Mutex + buf []byte + seq int64 + size int + full int + lastSize int + held int + closed bool + err error + + // ackMu is held while the watermark is read, which other uploads wait for + ackMu sync.Mutex + acked int64 + ackedAt time.Time + progress time.Time +} + +func newSlotWriter(ctx context.Context, store slotStore, table slotTable, p params) *slotWriter { + w := &slotWriter{ + ctx: ctx, + store: store, + table: table, + params: p, + slots: make(chan struct{}, len(table.Slots)), + size: min(slotInlineLimit, p.segmentBytes), + acked: -1, + progress: time.Now(), + } + go w.flushLoop() + return w +} + +func (w *slotWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + + if w.err != nil { + return 0, w.err + } + if w.closed { + return 0, io.ErrClosedPipe + } + + w.buf = append(w.buf, p...) + for len(w.buf) >= w.size { + if err := w.flushLocked(); err != nil { + return 0, err + } + } + return len(p), nil +} + +func (w *slotWriter) flushLoop() { + ticker := time.NewTicker(w.flushInterval) + defer ticker.Stop() + + for { + select { + case <-w.ctx.Done(): + return + case <-ticker.C: + w.mu.Lock() + if !w.closed && w.err == nil && len(w.buf) > 0 && w.readyToFlush() { + w.flushLocked() + } + done := w.closed || w.err != nil + w.mu.Unlock() + if done { + return + } + } + } +} + +func (w *slotWriter) readyToFlush() bool { + grew := len(w.buf) > w.lastSize + w.lastSize = len(w.buf) + + if grew && w.held < maxCoalescedTicks { + w.held++ + return false + } + w.held = 0 + return true +} + +// flushLocked sends a segment off buf once a slot is free, which it waits for +// without w.mu, so that a reader that falls behind holds Write back. w.mu is +// held +func (w *slotWriter) flushLocked() error { + w.mu.Unlock() + select { + case w.slots <- struct{}{}: + case <-w.ctx.Done(): + w.mu.Lock() + return w.ctx.Err() + } + w.mu.Lock() + if w.err != nil || len(w.buf) == 0 { + <-w.slots + return w.err + } + + n := min(len(w.buf), w.size) + if len(w.buf) > 4*slotInlineLimit { + // a burst goes at the full size at once + n = min(len(w.buf), w.segmentBytes) + } + if n == w.size && w.size < w.segmentBytes { + if w.full++; w.full == inlineRampSegments { + w.full = 0 + w.size = min(2*w.size, w.segmentBytes) + } + } + + chunk := bytes.Clone(w.buf[:n]) + w.buf = append(w.buf[:0], w.buf[n:]...) + seq := w.seq + w.seq++ + w.lastSize = len(w.buf) + w.held = 0 + + w.wg.Add(1) + go w.upload(seq, chunk) + return nil +} + +func (w *slotWriter) slotID(seq int64) string { + return w.table.Slots[seq%int64(len(w.table.Slots))] +} + +func (w *slotWriter) upload(seq int64, chunk []byte) { + defer w.wg.Done() + defer func() { <-w.slots }() + + mode := byte(slotMedia) + if len(chunk) <= slotInlineLimit { + mode = slotInline + } + if err := w.await(seq); err != nil { + w.fail(errors.New("failed to store segment ", seq).Base(err)) + return + } + if err := w.write(seq, mode, chunk); err != nil { + w.fail(errors.New("failed to store segment ", seq).Base(err)) + // the slot is free, the reader learns there what happened + w.write(seq, slotErr, nil) + } +} + +func (w *slotWriter) write(seq int64, mode byte, data []byte) error { + f := newSlotFrame(seq, mode, data) + var body []byte + if mode == slotMedia { + body = data + } + return w.store.setSlot(w.ctx, w.slotID(seq), f.encode(), body) +} + +// drained reports whether the reader took what the slot of seq held a lap +// before. The watermark is read again only after slotWatermarkTTL, as one +// that is old costs a wait but never what is in a slot +func (w *slotWriter) drained(seq int64) (bool, error) { + w.ackMu.Lock() + defer w.ackMu.Unlock() + + before := seq - int64(len(w.table.Slots)) + if w.acked >= before { + return true, nil + } + if time.Since(w.ackedAt) < slotWatermarkTTL { + return false, nil + } + acked, err := readWatermark(w.ctx, w.store, w.table.Watermark) + if err != nil { + return false, err + } + w.ackedAt = time.Now() + if acked > w.acked { + w.acked, w.progress = acked, w.ackedAt + } + return w.acked >= before, nil +} + +// await waits for drained, for slotStallTimeout at most, or slotCloseGrace +// without progress once the writer is closed +func (w *slotWriter) await(seq int64) error { + start := time.Now() + wait := slotAckRetry + for { + ok, err := w.drained(seq) + if ok { + return nil + } + w.mu.Lock() + closed := w.closed + w.mu.Unlock() + w.ackMu.Lock() + still := time.Since(w.progress) + w.ackMu.Unlock() + if time.Since(start) > slotStallTimeout || closed && still > slotCloseGrace { + return errors.New("the reader took nothing from slot ", w.slotID(seq), " for ", still).Base(err) + } + + if err != nil { + wait = min(2*wait, 2*time.Second) + } else { + wait = slotAckRetry + } + select { + case <-w.ctx.Done(): + return w.ctx.Err() + case <-time.After(wait): + } + } +} + +func (w *slotWriter) fail(err error) { + w.mu.Lock() + defer w.mu.Unlock() + if w.err == nil { + w.err = err + } +} + +func (w *slotWriter) Close() error { + w.mu.Lock() + if w.closed { + w.mu.Unlock() + return nil + } + w.closed = true + for len(w.buf) > 0 && w.err == nil { + if err := w.flushLocked(); err != nil { + w.mu.Unlock() + return err + } + } + seq := w.seq + w.mu.Unlock() + + w.ackMu.Lock() + w.progress = time.Now() + w.ackMu.Unlock() + + w.wg.Wait() + + w.mu.Lock() + err := w.err + w.mu.Unlock() + if err != nil { + return err + } + if err := w.await(seq); err != nil { + return err + } + return w.write(seq, slotEnd, nil) +} + +type slotReader struct { + ctx context.Context + store slotStore + table slotTable + params + + ch chan []byte + wake chan struct{} + acks chan struct{} + + mu sync.Mutex + head int64 + started bool + // hitAt is when a segment last came and gap the time between them, activeAt + // when a segment came, was seen coming or Wake. Nothing for the eager window + // is idle, and active is closed when that ends + hitAt time.Time + gap time.Duration + activeAt time.Time + active chan struct{} + err error +} + +func newSlotReader(ctx context.Context, store slotStore, table slotTable, p params) *slotReader { + r := &slotReader{ + ctx: ctx, + store: store, + table: table, + params: p, + ch: make(chan []byte, p.concurrency), + wake: make(chan struct{}, 1), + acks: make(chan struct{}, 1), + hitAt: time.Now(), + activeAt: time.Now(), + active: make(chan struct{}), + } + go r.run() + go r.ackLoop() + return r +} + +func (r *slotReader) chunks() <-chan []byte { + return r.ch +} + +// Wake has the head asked for now, as an answer is likely to come +func (r *slotReader) Wake() { + r.mu.Lock() + r.touch() + r.mu.Unlock() + select { + case r.wake <- struct{}{}: + default: + } +} + +func (r *slotReader) slotID(seq int64) string { + return r.table.Slots[seq%int64(len(r.table.Slots))] +} + +type slotFetch struct { + ready chan struct{} + // promoted is closed when the segment becomes the head + promoted chan struct{} + body []byte + end bool + err error +} + +func (r *slotReader) run() { + defer close(r.ch) + ctx, cancel := context.WithCancel(r.ctx) + defer cancel() + + fetches := make(map[int64]*slotFetch) + for { + r.mu.Lock() + head, started := r.head, r.started + r.mu.Unlock() + + window := 1 + if started { + window = min(slotFetchAhead, len(r.table.Slots)) + } + for i := range window { + seq := head + int64(i) + if fetches[seq] == nil { + fetches[seq] = &slotFetch{ready: make(chan struct{}), promoted: make(chan struct{})} + go r.fetch(ctx, seq, i, fetches[seq]) + } + } + + f := fetches[head] + select { + case <-f.ready: + case <-ctx.Done(): + return + } + delete(fetches, head) + if f.err != nil { + r.setErr(f.err) + return + } + if f.end { + return + } + + // the slot is free once the bytes are here, not once they are read + r.mu.Lock() + r.head++ + r.started = true + r.mu.Unlock() + if next := fetches[head+1]; next != nil { + close(next.promoted) + } + select { + case r.acks <- struct{}{}: + default: + } + + select { + case r.ch <- f.body: + case <-ctx.Done(): + return + } + } +} + +func (r *slotReader) isHead(seq int64) bool { + r.mu.Lock() + defer r.mu.Unlock() + return r.head == seq +} + +// fetch asks for the slot of seq until it holds seq: the head often, a segment +// behind it at the pace that segments come, as it would wait for the head anyway +func (r *slotReader) fetch(ctx context.Context, seq int64, offset int, f *slotFetch) { + defer close(f.ready) + if offset > 0 && !r.hold(ctx, time.Duration(offset)*r.pace(), f.promoted, nil) { + return + } + + id := r.slotID(seq) + backoff := slotHuntFloor + var goneSince, warnedAt time.Time + for { + description, err := r.store.readDescription(ctx, id) + frame, ok := decodeSlot(description) + switch { + case err == nil && ok && frame.seq == seq: + // what follows it is asked for while its body comes + r.mu.Lock() + r.touch() + r.mu.Unlock() + switch frame.mode { + case slotEnd: + f.end = true + return + case slotErr: + f.err = errors.New("the peer could not store segment ", seq) + return + case slotInline: + f.body = frame.data + r.hit() + return + case slotMedia: + if body, err := r.store.readBody(ctx, id); err == nil && frame.holds(body) { + f.body = body + r.hit() + return + } + } + case err == errNotFound: + if goneSince.IsZero() { + goneSince = time.Now() + } else if time.Since(goneSince) > slotGoneTimeout && r.isHead(seq) { + f.err = errors.New("slot ", id, " of segment ", seq, " is gone").Base(err) + return + } + case err != nil: + goneSince = time.Time{} + if ctx.Err() == nil && time.Since(warnedAt) > 10*time.Second { + warnedAt = time.Now() + errors.LogWarningInner(ctx, err, "failed to read slot ", id) + } + } + + r.mu.Lock() + head, idle, active := r.head == seq, time.Since(r.activeAt) > r.eagerWindow, r.active + r.mu.Unlock() + wait, a, b := backoff, (<-chan struct{})(r.wake), (<-chan struct{})(nil) + switch { + case head: + ceiling := slotHuntCeiling + if idle { + ceiling = max(ceiling, r.maxPollInterval) + } + backoff = min(2*backoff, ceiling) + case idle: + // only the head is asked for until something comes + wait, a, b, backoff = 0, f.promoted, active, slotHuntFloor + default: + wait, a, backoff = min(max(r.pace(), slotTailMin), slotTailMax), f.promoted, slotHuntFloor + } + if !r.hold(ctx, wait, a, b) { + return + } + } +} + +// pace is about how long segments take to come one after the other +func (r *slotReader) pace() time.Duration { + r.mu.Lock() + defer r.mu.Unlock() + return min(max(r.gap, 300*time.Millisecond), 3*time.Second) +} + +func (r *slotReader) hit() { + r.mu.Lock() + defer r.mu.Unlock() + if gap := time.Since(r.hitAt); r.gap == 0 { + r.gap = gap + } else { + r.gap = (r.gap + gap) / 2 + } + r.hitAt = time.Now() + r.touch() +} + +// touch notes that something came or is likely to. r.mu is held +func (r *slotReader) touch() { + if time.Since(r.activeAt) > r.eagerWindow { + close(r.active) + r.active = make(chan struct{}) + } + r.activeAt = time.Now() +} + +// hold waits d, unless d is 0, or until a or b is closed or has a value +func (r *slotReader) hold(ctx context.Context, d time.Duration, a, b <-chan struct{}) bool { + var timeout <-chan time.Time + if d > 0 { + timer := time.NewTimer(d) + defer timer.Stop() + timeout = timer.C + } + select { + case <-ctx.Done(): + return false + case <-a: + case <-b: + case <-timeout: + } + return true +} + +// ackLoop writes the watermark when the head moves, the last head only if it +// moved more than once meanwhile +func (r *slotReader) ackLoop() { + acked := int64(-1) + wait := slotAckRetry + for { + select { + case <-r.ctx.Done(): + return + case <-r.acks: + } + for { + r.mu.Lock() + want := r.head - 1 + r.mu.Unlock() + if want <= acked { + break + } + if err := r.store.setSlot(r.ctx, r.table.Watermark, strconv.FormatInt(want, 10), nil); err == nil { + acked, wait = want, slotAckRetry + continue + } + wait = min(2*wait, 5*time.Second) + select { + case <-r.ctx.Done(): + return + case <-time.After(wait): + } + } + } +} + +func (r *slotReader) setErr(err error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.err == nil { + r.err = err + } +} + +func (r *slotReader) Err() error { + r.mu.Lock() + defer r.mu.Unlock() + if r.err != nil { + return r.err + } + return io.EOF +} diff --git a/transport/internet/xdrive/slots_test.go b/transport/internet/xdrive/slots_test.go new file mode 100644 index 000000000000..1a77e873ecaa --- /dev/null +++ b/transport/internet/xdrive/slots_test.go @@ -0,0 +1,307 @@ +package xdrive + +import ( + "bytes" + "context" + "encoding/base64" + "errors" + "io" + "math/rand" + "os" + "strconv" + "sync" + "testing" + "time" +) + +type memSlots struct { + mu sync.Mutex + next int + descs map[string]string + bodies map[string][]byte + reads int +} + +func newMemSlots() *memSlots { + return &memSlots{descs: make(map[string]string), bodies: make(map[string][]byte)} +} + +func (m *memSlots) makeSlot(ctx context.Context, name string) (string, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.next++ + id := "slot-" + strconv.Itoa(m.next) + m.descs[id] = "" + return id, nil +} + +func (m *memSlots) setSlot(ctx context.Context, id, description string, body []byte) error { + m.mu.Lock() + defer m.mu.Unlock() + if _, ok := m.descs[id]; !ok { + return errNotFound + } + m.descs[id] = description + if body != nil { + m.bodies[id] = bytes.Clone(body) + } + return nil +} + +func (m *memSlots) readDescription(ctx context.Context, id string) (string, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.reads++ + d, ok := m.descs[id] + if !ok { + return "", errNotFound + } + return d, nil +} + +func (m *memSlots) readBody(ctx context.Context, id string) ([]byte, error) { + m.mu.Lock() + defer m.mu.Unlock() + return m.bodies[id], nil +} + +// the slots of all tests are asked for often, set once as a goroutine of a +// closed Conn may still read them +func init() { + slotHuntFloor, slotHuntCeiling, slotTailMin, slotTailMax = 2*time.Millisecond, 10*time.Millisecond, 5*time.Millisecond, 20*time.Millisecond + slotAckRetry, slotWatermarkTTL = 2*time.Millisecond, 2*time.Millisecond +} + +func slotPair(t *testing.T, store slotStore, concurrency, segmentBytes uint32) (*Conn, *Conn) { + t.Helper() + p := paramsFromConfig(&Config{ + SegmentBytes: segmentBytes, + FlushIntervalMs: 1, + PollIntervalMs: 1, + MaxPollIntervalMs: 20, + EagerWindowMs: 50, + Concurrency: concurrency, + }) + r, err := makeRing(context.Background(), store, "session", p) + if err != nil { + t.Fatal(err) + } + client := newSlotConn(context.Background(), store, r.Up, r.Down, p, nil) + server := newSlotConn(context.Background(), store, r.Down, r.Up, p, nil) + t.Cleanup(func() { + client.Close() + server.Close() + }) + for _, c := range []*Conn{client, server} { + c.SetReadDeadline(time.Now().Add(30 * time.Second)) + } + return client, server +} + +// writeAll writes payload in pieces of random size, and closes c if end, the +// error goes to errs +func writeAll(c *Conn, payload []byte, end bool, errs chan<- error) { + go func() { + rng := rand.New(rand.NewSource(1)) + for len(payload) > 0 { + n := min(len(payload), 1+rng.Intn(40<<10)) + if _, err := c.Write(payload[:n]); err != nil { + errs <- err + return + } + payload = payload[n:] + } + var err error + if end { + err = c.Close() + } + errs <- err + }() +} + +func readAll(t *testing.T, c *Conn, want []byte) { + t.Helper() + got, err := io.ReadAll(c) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, want) { + t.Fatalf("got %d bytes, not the %d that were sent", len(got), len(want)) + } +} + +func randomPayload(n int) []byte { + payload := make([]byte, n) + rand.New(rand.NewSource(int64(n))).Read(payload) + return payload +} + +func TestSlotFrame(t *testing.T) { + data := []byte("hello") + f, ok := decodeSlot(newSlotFrame(7, slotInline, data).encode()) + if !ok || f.seq != 7 || f.mode != slotInline || !bytes.Equal(f.data, data) { + t.Fatalf("inline frame came back as %+v %v", f, ok) + } + + body := randomPayload(100) + f, ok = decodeSlot(newSlotFrame(8, slotMedia, body).encode()) + if !ok || f.seq != 8 || f.mode != slotMedia || f.data != nil || !f.holds(body) || f.holds(body[1:]) { + t.Fatalf("media frame came back as %+v %v", f, ok) + } + + corrupt, _ := base64.StdEncoding.DecodeString(newSlotFrame(1, slotInline, data).encode()) + corrupt[len(corrupt)-1]++ + for _, broken := range []string{"", "not base64", newSlotFrame(1, slotInline, data).encode()[:20], "WERSVlNMMDE=", + base64.StdEncoding.EncodeToString(corrupt)} { + if _, ok := decodeSlot(broken); ok { + t.Errorf("%q decoded", broken) + } + } +} + +// both sides send more segments than the ring has slots at once, small ones +// that go in descriptions and large ones that go in bodies, then one closes +// and the other reads EOF +func TestSlotConnRoundTrip(t *testing.T) { + client, server := slotPair(t, newMemSlots(), 4, 64<<10) + up, down := randomPayload(1<<20), randomPayload(300<<10) + errs := make(chan error, 2) + writeAll(client, up, false, errs) + writeAll(server, down, false, errs) + + for _, c := range []struct { + conn *Conn + want []byte + }{{server, up}, {client, down}} { + got := make([]byte, len(c.want)) + if _, err := io.ReadFull(c.conn, got); err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, c.want) { + t.Fatal("what came is not what was sent") + } + } + for range 2 { + if err := <-errs; err != nil { + t.Fatal(err) + } + } + + if err := client.Close(); err != nil { + t.Fatal(err) + } + readAll(t, server, nil) +} + +// a reader that does not read holds the writer back instead of losing what a +// lap of the ring would overwrite +func TestSlotConnSlowReader(t *testing.T) { + client, server := slotPair(t, newMemSlots(), 2, 16<<10) + payload := randomPayload(400 << 10) + errs := make(chan error, 1) + writeAll(client, payload, true, errs) + + time.Sleep(200 * time.Millisecond) + readAll(t, server, payload) + if err := <-errs; err != nil { + t.Fatal(err) + } +} + +// the next segment is asked for at once when it becomes the head, not after +// the wait it had as one that was fetched ahead +func TestSlotLatency(t *testing.T) { + client, server := slotPair(t, newMemSlots(), 8, 64<<10) + start := time.Now() + for range 10 { + for _, c := range [][2]*Conn{{client, server}, {server, client}} { + c[0].Write([]byte("ping")) + if _, err := io.ReadFull(c[1], make([]byte, 4)); err != nil { + t.Fatal(err) + } + } + } + if took := time.Since(start); took > 2*time.Second { + t.Fatalf("10 round trips took %v", took) + } +} + +// once nothing comes for the eager window only the heads are asked for, not +// the segments that a reader fetches ahead while things come +func TestSlotIdle(t *testing.T) { + store := newMemSlots() + client, server := slotPair(t, store, 8, 64<<10) + for _, c := range [][2]*Conn{{client, server}, {server, client}} { + c[0].Write([]byte("ping")) + if _, err := io.ReadFull(c[1], make([]byte, 4)); err != nil { + t.Fatal(err) + } + } + + time.Sleep(200 * time.Millisecond) + store.mu.Lock() + before := store.reads + store.mu.Unlock() + time.Sleep(500 * time.Millisecond) + store.mu.Lock() + reads := store.reads - before + store.mu.Unlock() + // two heads, each asked for every 20 ms at most + if reads > 75 { + t.Fatalf("an idle session asked for slots %d times in 500 ms", reads) + } +} + +type failingSlots struct { + *memSlots +} + +// setSlot fails for segments, not for the error frame +func (f failingSlots) setSlot(ctx context.Context, id, description string, body []byte) error { + if frame, ok := decodeSlot(description); ok && (frame.mode == slotInline || frame.mode == slotMedia) { + return io.ErrUnexpectedEOF + } + return f.memSlots.setSlot(ctx, id, description, body) +} + +func TestSlotWriterFailure(t *testing.T) { + client, server := slotPair(t, failingSlots{newMemSlots()}, 2, 16<<10) + client.Write([]byte("lost")) + server.SetReadDeadline(time.Now().Add(5 * time.Second)) + if _, err := server.Read(make([]byte, 10)); err == nil || err == io.EOF || errors.Is(err, os.ErrDeadlineExceeded) { + t.Fatalf("the reader got %v, not the failure of the writer", err) + } + for start := time.Now(); time.Since(start) < 5*time.Second; time.Sleep(time.Millisecond) { + if _, err := client.Write([]byte("more")); err != nil { + return + } + } + t.Fatal("Write did not fail") +} + +type staleSlots struct { + *memSlots + mu sync.Mutex + served map[string]bool +} + +// readBody gives what a slot held before once, as a body read right after its +// description may +func (s *staleSlots) readBody(ctx context.Context, id string) ([]byte, error) { + s.mu.Lock() + first := !s.served[id] + s.served[id] = true + s.mu.Unlock() + if first { + return []byte("stale"), nil + } + return s.memSlots.readBody(ctx, id) +} + +func TestSlotStaleBody(t *testing.T) { + client, server := slotPair(t, &staleSlots{memSlots: newMemSlots(), served: make(map[string]bool)}, 2, 64<<10) + payload := randomPayload(200 << 10) + errs := make(chan error, 1) + writeAll(client, payload, true, errs) + readAll(t, server, payload) +} diff --git a/transport/internet/xdrive/wal.go b/transport/internet/xdrive/wal.go index 62e1876f2638..a05def440eb1 100644 --- a/transport/internet/xdrive/wal.go +++ b/transport/internet/xdrive/wal.go @@ -230,6 +230,10 @@ func newWALReader(ctx context.Context, storage Storage, prefix string, p params) return r } +func (r *walReader) chunks() <-chan []byte { + return r.ch +} + func (r *walReader) Wake() { select { case r.wake <- struct{}{}: diff --git a/transport/internet/xdrive/xdrive.go b/transport/internet/xdrive/xdrive.go index e79e738d32e2..c326d506c5f4 100644 --- a/transport/internet/xdrive/xdrive.go +++ b/transport/internet/xdrive/xdrive.go @@ -4,6 +4,7 @@ import ( "context" "crypto/rand" "encoding/hex" + "encoding/json" "fmt" "strconv" "strings" @@ -95,15 +96,33 @@ func Dial(ctx context.Context, dest net.Destination, streamSettings *internet.Me return nil, err } - if err := storage.Put(ctx, announceName(session, time.Now()), nil); err != nil { + p := paramsFromConfig(config) + var announce []byte + var r ring + slots, ok := storage.(slotStore) + if ok { + if r, err = makeRing(ctx, slots, session, p); err == nil { + announce, _ = json.Marshal(r) + } else { + ok = false + errors.LogWarningInner(ctx, err, "failed to make the slots of session ", session, ", it goes in segments") + } + } + + if err := storage.Put(ctx, announceName(session, time.Now()), announce); err != nil { storage.Close() return nil, errors.New("failed to announce session ", session).Base(err) } + if ok { + errors.LogInfo(ctx, "opened session ", session, " in ", len(r.Up.Slots), " slots") + return newSlotConn(context.Background(), slots, r.Up, r.Down, p, func() { + storage.Close() + }), nil + } errors.LogInfo(ctx, "opened session ", session) - return newConn(context.Background(), storage, - uplinkPrefix(session), downlinkPrefix(session), paramsFromConfig(config), func() { + uplinkPrefix(session), downlinkPrefix(session), p, func() { storage.Close() }), nil } @@ -206,9 +225,8 @@ func (l *Listener) acceptPending(logCtx context.Context) (bool, error) { continue } go l.drop(full) - errors.LogInfo(logCtx, "accepted session ", session) accepted = true - l.addConn(l.newSessionConn(session)) + l.addConn(l.newSessionConn(logCtx, session, listed.Inline)) } return accepted, nil } @@ -228,15 +246,24 @@ func (l *Listener) claim(session string) bool { return true } -func (l *Listener) newSessionConn(session string) *Conn { +// newSessionConn takes a session in the slots that its announcement has, or +// in segments if it has none +func (l *Listener) newSessionConn(logCtx context.Context, session string, announce []byte) *Conn { storage := sessionStorage(l.storage) - return newConn(l.ctx, storage, - downlinkPrefix(session), uplinkPrefix(session), l.params, func() { - storage.Close() - l.mu.Lock() - delete(l.active, session) - l.mu.Unlock() - }) + onClose := func() { + storage.Close() + l.mu.Lock() + delete(l.active, session) + l.mu.Unlock() + } + if slots, ok := storage.(slotStore); ok { + if r, ok := decodeRing(announce); ok { + errors.LogInfo(logCtx, "accepted session ", session, " in ", len(r.Down.Slots), " slots") + return newSlotConn(l.ctx, slots, r.Down, r.Up, l.params, onClose) + } + } + errors.LogInfo(logCtx, "accepted session ", session) + return newConn(l.ctx, storage, downlinkPrefix(session), uplinkPrefix(session), l.params, onClose) } func (l *Listener) collectLoop(logCtx context.Context) {