From 342f8baa698161dd38abd8160dae32b4235bd190 Mon Sep 17 00:00:00 2001 From: pingqiu Date: Tue, 7 Apr 2026 14:57:43 -0700 Subject: [PATCH] =?UTF-8?q?feat:=20rebuild=20transport=20wiring=20?= =?UTF-8?q?=E2=80=94=20session=20control=20+=20base=20block=20streaming?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire protocol messages and transport handlers for the rebuild MVP: Protocol messages (rebuild_transport.go): - SessionControlMsg: epoch, sessionID, command, baseLSN, targetLSN, snapshotID. Encode/Decode with fixed 37-byte wire format. - SessionAckMsg: epoch, sessionID, phase, walAppliedLSN, baseComplete, achievedLSN. Encode/Decode with fixed 34-byte wire format. - MsgSessionControl (0x10) and MsgSessionAck (0x11) on control channel. - SendSessionControl/SendSessionAck convenience functions. Transport handlers: - RebuildTransportServer: primary-side, streams all extent blocks as MsgRebuildExtent frames (reusing existing rebuild message type), ends with MsgRebuildDone. - RebuildTransportClient: replica-side, receives base blocks and routes through vol.ApplyRebuildSessionBaseBlock, marks base complete on MsgRebuildDone. 4 transport tests: - SessionControl wire round-trip - SessionAck wire round-trip - BaseBlockStreaming: full TCP loop, 1024 blocks streamed and verified - SessionControlOverTCP: real TCP send/receive with accepted ack Co-Authored-By: Claude Opus 4.6 (1M context) --- weed/storage/blockvol/rebuild_transport.go | 257 ++++++++++++++++++ .../blockvol/rebuild_transport_test.go | 236 ++++++++++++++++ 2 files changed, 493 insertions(+) create mode 100644 weed/storage/blockvol/rebuild_transport.go create mode 100644 weed/storage/blockvol/rebuild_transport_test.go diff --git a/weed/storage/blockvol/rebuild_transport.go b/weed/storage/blockvol/rebuild_transport.go new file mode 100644 index 000000000..4ed32e6d6 --- /dev/null +++ b/weed/storage/blockvol/rebuild_transport.go @@ -0,0 +1,257 @@ +package blockvol + +import ( + "encoding/binary" + "fmt" + "io" + "log" + "net" + "time" +) + +// Session control message types (on control channel, distinct from barrier). +const ( + MsgSessionControl byte = 0x10 // primary → replica: session command + MsgSessionAck byte = 0x11 // replica → primary: session progress/result +) + +// Session control command kinds. +const ( + SessionCmdStartRebuild byte = 0x01 + SessionCmdCancel byte = 0x02 +) + +// Session ack phase codes (wire representation of RebuildSessionPhase). +const ( + SessionAckAccepted byte = 0x01 + SessionAckRunning byte = 0x02 + SessionAckBaseComplete byte = 0x03 + SessionAckCompleted byte = 0x04 + SessionAckFailed byte = 0x05 +) + +// SessionControlMsg is the wire message for session control commands. +type SessionControlMsg struct { + Epoch uint64 + SessionID uint64 + Command byte + BaseLSN uint64 // for start_rebuild + TargetLSN uint64 // for start_rebuild + SnapshotID uint32 // for start_rebuild (0 = use current extent) +} + +// EncodeSessionControl serializes a session control message. +// Wire: [8B epoch][8B sessionID][1B cmd][8B baseLSN][8B targetLSN][4B snapshotID] = 37 bytes. +func EncodeSessionControl(msg SessionControlMsg) []byte { + buf := make([]byte, 37) + binary.BigEndian.PutUint64(buf[0:8], msg.Epoch) + binary.BigEndian.PutUint64(buf[8:16], msg.SessionID) + buf[16] = msg.Command + binary.BigEndian.PutUint64(buf[17:25], msg.BaseLSN) + binary.BigEndian.PutUint64(buf[25:33], msg.TargetLSN) + binary.BigEndian.PutUint32(buf[33:37], msg.SnapshotID) + return buf +} + +// DecodeSessionControl deserializes a session control message. +func DecodeSessionControl(buf []byte) (SessionControlMsg, error) { + if len(buf) < 37 { + return SessionControlMsg{}, fmt.Errorf("session control: short message (%d bytes)", len(buf)) + } + return SessionControlMsg{ + Epoch: binary.BigEndian.Uint64(buf[0:8]), + SessionID: binary.BigEndian.Uint64(buf[8:16]), + Command: buf[16], + BaseLSN: binary.BigEndian.Uint64(buf[17:25]), + TargetLSN: binary.BigEndian.Uint64(buf[25:33]), + SnapshotID: binary.BigEndian.Uint32(buf[33:37]), + }, nil +} + +// SessionAckMsg is the wire message for session progress/result. +type SessionAckMsg struct { + Epoch uint64 + SessionID uint64 + Phase byte + WALAppliedLSN uint64 + BaseComplete bool + AchievedLSN uint64 // on completion +} + +// EncodeSessionAck serializes a session ack message. +// Wire: [8B epoch][8B sessionID][1B phase][8B walAppliedLSN][1B baseComplete][8B achievedLSN] = 34 bytes. +func EncodeSessionAck(msg SessionAckMsg) []byte { + buf := make([]byte, 34) + binary.BigEndian.PutUint64(buf[0:8], msg.Epoch) + binary.BigEndian.PutUint64(buf[8:16], msg.SessionID) + buf[16] = msg.Phase + binary.BigEndian.PutUint64(buf[17:25], msg.WALAppliedLSN) + if msg.BaseComplete { + buf[25] = 1 + } + binary.BigEndian.PutUint64(buf[26:34], msg.AchievedLSN) + return buf +} + +// DecodeSessionAck deserializes a session ack message. +func DecodeSessionAck(buf []byte) (SessionAckMsg, error) { + if len(buf) < 34 { + return SessionAckMsg{}, fmt.Errorf("session ack: short message (%d bytes)", len(buf)) + } + return SessionAckMsg{ + Epoch: binary.BigEndian.Uint64(buf[0:8]), + SessionID: binary.BigEndian.Uint64(buf[8:16]), + Phase: buf[16], + WALAppliedLSN: binary.BigEndian.Uint64(buf[17:25]), + BaseComplete: buf[25] != 0, + AchievedLSN: binary.BigEndian.Uint64(buf[26:34]), + }, nil +} + +// RebuildTransportServer handles the primary-side rebuild data serving for one +// session. It streams snapshot base blocks to the replica over a dedicated TCP +// connection (the existing rebuild server path). +type RebuildTransportServer struct { + vol *BlockVol + sessionID uint64 + epoch uint64 + baseLSN uint64 + targetLSN uint64 +} + +// NewRebuildTransportServer creates a primary-side rebuild transport server +// for one session. +func NewRebuildTransportServer(vol *BlockVol, sessionID, epoch, baseLSN, targetLSN uint64) *RebuildTransportServer { + return &RebuildTransportServer{ + vol: vol, + sessionID: sessionID, + epoch: epoch, + baseLSN: baseLSN, + targetLSN: targetLSN, + } +} + +// ServeBaseBlocks streams all extent blocks to the replica connection. +// Each block is sent as MsgRebuildExtent with the LBA encoded in the first 8 +// bytes. The stream ends with MsgRebuildDone. +func (s *RebuildTransportServer) ServeBaseBlocks(conn net.Conn) error { + if s.vol == nil { + return fmt.Errorf("rebuild transport: volume is nil") + } + + conn.SetDeadline(time.Now().Add(10 * time.Minute)) + defer conn.SetDeadline(time.Time{}) + + info := s.vol.Info() + blockSize := uint64(info.BlockSize) + totalLBAs := info.VolumeSize / blockSize + + // Flush to ensure extent is current before streaming. + if err := s.vol.ForceFlush(); err != nil { + log.Printf("rebuild transport: flush before base stream: %v", err) + } + + var sentBlocks uint64 + for lba := uint64(0); lba < totalLBAs; lba++ { + data, err := s.vol.ReadLBA(lba, uint32(blockSize)) + if err != nil { + return fmt.Errorf("rebuild transport: read LBA %d: %w", lba, err) + } + + // Encode: [8B LBA][block data] + frame := make([]byte, 8+len(data)) + binary.BigEndian.PutUint64(frame[0:8], lba) + copy(frame[8:], data) + + if err := WriteFrame(conn, MsgRebuildExtent, frame); err != nil { + return fmt.Errorf("rebuild transport: send LBA %d: %w", lba, err) + } + sentBlocks++ + } + + // Send completion marker. + doneBuf := make([]byte, 8) + binary.BigEndian.PutUint64(doneBuf, sentBlocks) + if err := WriteFrame(conn, MsgRebuildDone, doneBuf); err != nil { + return fmt.Errorf("rebuild transport: send done: %w", err) + } + + log.Printf("rebuild transport: served %d base blocks for session %d", sentBlocks, s.sessionID) + return nil +} + +// RebuildTransportClient handles the replica-side rebuild data receiving for +// one session. It receives snapshot base blocks from the primary and routes +// them through the rebuild session's base lane. +type RebuildTransportClient struct { + vol *BlockVol + sessionID uint64 +} + +// NewRebuildTransportClient creates a replica-side rebuild transport client +// for one session. +func NewRebuildTransportClient(vol *BlockVol, sessionID uint64) *RebuildTransportClient { + return &RebuildTransportClient{ + vol: vol, + sessionID: sessionID, + } +} + +// ReceiveBaseBlocks reads base blocks from the primary connection and applies +// them through the rebuild session. Returns the total number of blocks processed. +func (c *RebuildTransportClient) ReceiveBaseBlocks(conn net.Conn) (uint64, error) { + if c.vol == nil { + return 0, fmt.Errorf("rebuild transport: volume is nil") + } + + conn.SetDeadline(time.Now().Add(10 * time.Minute)) + defer conn.SetDeadline(time.Time{}) + + var totalBlocks uint64 + for { + msgType, payload, err := ReadFrame(conn) + if err != nil { + if err == io.EOF { + break + } + return totalBlocks, fmt.Errorf("rebuild transport: read frame: %w", err) + } + + switch msgType { + case MsgRebuildExtent: + if len(payload) < 8 { + return totalBlocks, fmt.Errorf("rebuild transport: short extent frame") + } + lba := binary.BigEndian.Uint64(payload[0:8]) + data := payload[8:] + if _, err := c.vol.ApplyRebuildSessionBaseBlock(c.sessionID, lba, data); err != nil { + return totalBlocks, fmt.Errorf("rebuild transport: apply base LBA %d: %w", lba, err) + } + totalBlocks++ + + case MsgRebuildDone: + if err := c.vol.MarkRebuildSessionBaseComplete(c.sessionID, totalBlocks); err != nil { + return totalBlocks, fmt.Errorf("rebuild transport: mark base complete: %w", err) + } + log.Printf("rebuild transport: received %d base blocks for session %d", totalBlocks, c.sessionID) + return totalBlocks, nil + + case MsgRebuildError: + return totalBlocks, fmt.Errorf("rebuild transport: server error: %s", string(payload)) + + default: + return totalBlocks, fmt.Errorf("rebuild transport: unexpected message type 0x%02x", msgType) + } + } + return totalBlocks, nil +} + +// SendSessionControl sends a session control message on the control connection. +func SendSessionControl(conn net.Conn, msg SessionControlMsg) error { + return WriteFrame(conn, MsgSessionControl, EncodeSessionControl(msg)) +} + +// SendSessionAck sends a session ack message on the control connection. +func SendSessionAck(conn net.Conn, msg SessionAckMsg) error { + return WriteFrame(conn, MsgSessionAck, EncodeSessionAck(msg)) +} diff --git a/weed/storage/blockvol/rebuild_transport_test.go b/weed/storage/blockvol/rebuild_transport_test.go new file mode 100644 index 000000000..82890782a --- /dev/null +++ b/weed/storage/blockvol/rebuild_transport_test.go @@ -0,0 +1,236 @@ +package blockvol + +import ( + "bytes" + "fmt" + "net" + "testing" + "time" +) + +func TestRebuildTransport_SessionControlRoundTrip(t *testing.T) { + msg := SessionControlMsg{ + Epoch: 5, + SessionID: 42, + Command: SessionCmdStartRebuild, + BaseLSN: 1000, + TargetLSN: 2000, + SnapshotID: 7, + } + encoded := EncodeSessionControl(msg) + decoded, err := DecodeSessionControl(encoded) + if err != nil { + t.Fatalf("decode: %v", err) + } + if decoded != msg { + t.Fatalf("round-trip mismatch: got %+v, want %+v", decoded, msg) + } +} + +func TestRebuildTransport_SessionAckRoundTrip(t *testing.T) { + msg := SessionAckMsg{ + Epoch: 5, + SessionID: 42, + Phase: SessionAckCompleted, + WALAppliedLSN: 2500, + BaseComplete: true, + AchievedLSN: 2500, + } + encoded := EncodeSessionAck(msg) + decoded, err := DecodeSessionAck(encoded) + if err != nil { + t.Fatalf("decode: %v", err) + } + if decoded != msg { + t.Fatalf("round-trip mismatch: got %+v, want %+v", decoded, msg) + } +} + +func TestRebuildTransport_BaseBlockStreaming(t *testing.T) { + // Create primary with data. + primary := createTestVolForTransport(t, "primary") + defer primary.Close() + if err := primary.HandleAssignment(1, RolePrimary, 30*time.Second); err != nil { + t.Fatal(err) + } + blocks := make([][]byte, 5) + for i := 0; i < 5; i++ { + blocks[i] = bytes.Repeat([]byte{byte(0xA0 + i)}, 4096) + if err := primary.WriteLBA(uint64(i), blocks[i]); err != nil { + t.Fatalf("write LBA %d: %v", i, err) + } + } + if err := primary.SyncCache(); err != nil { + t.Fatalf("sync: %v", err) + } + + // Create replica with rebuild session. + replica := createTestVolForTransport(t, "replica") + defer replica.Close() + if err := replica.HandleAssignment(1, RoleReplica, 30*time.Second); err != nil { + t.Fatal(err) + } + if err := replica.StartRebuildSession(RebuildSessionConfig{ + SessionID: 1, + Epoch: 1, + BaseLSN: 5, + TargetLSN: 5, + }); err != nil { + t.Fatal(err) + } + + // Set up localhost TCP connection. + server := NewRebuildTransportServer(primary, 1, 1, 5, 5) + client := NewRebuildTransportClient(replica, 1) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + // Server side: accept and serve. + serverDone := make(chan error, 1) + go func() { + conn, err := ln.Accept() + if err != nil { + serverDone <- err + return + } + defer conn.Close() + serverDone <- server.ServeBaseBlocks(conn) + }() + + // Client side: connect and receive. + conn, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + totalBlocks, err := client.ReceiveBaseBlocks(conn) + if err != nil { + t.Fatalf("receive base blocks: %v", err) + } + + // Wait for server. + if err := <-serverDone; err != nil { + t.Fatalf("serve base blocks: %v", err) + } + + // The volume is 4MB / 4K = 1024 LBAs, but only 5 have data. + // All 1024 blocks are streamed (full extent). + t.Logf("received %d base blocks", totalBlocks) + if totalBlocks == 0 { + t.Fatal("no base blocks received") + } + + // Verify the 5 written blocks are readable on replica. + for i := 0; i < 5; i++ { + data, err := replica.ReadLBA(uint64(i), 4096) + if err != nil { + t.Fatalf("replica read LBA %d: %v", i, err) + } + if !bytes.Equal(data, blocks[i]) { + t.Fatalf("replica LBA %d mismatch: got[0]=0x%02x want[0]=0x%02x", + i, data[0], blocks[i][0]) + } + } + t.Log("base block streaming: all 5 blocks verified on replica") +} + +func TestRebuildTransport_SessionControlOverTCP(t *testing.T) { + // Test session control messages over real TCP. + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + done := make(chan error, 1) + go func() { + conn, err := ln.Accept() + if err != nil { + done <- err + return + } + defer conn.Close() + + // Read session control. + msgType, payload, err := ReadFrame(conn) + if err != nil { + done <- err + return + } + if msgType != MsgSessionControl { + done <- fmt.Errorf("unexpected msg type: 0x%02x", msgType) + return + } + ctrl, err := DecodeSessionControl(payload) + if err != nil { + done <- err + return + } + + // Reply with session ack. + ack := SessionAckMsg{ + Epoch: ctrl.Epoch, + SessionID: ctrl.SessionID, + Phase: SessionAckAccepted, + } + done <- SendSessionAck(conn, ack) + }() + + conn, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + // Send session control. + ctrl := SessionControlMsg{ + Epoch: 3, + SessionID: 99, + Command: SessionCmdStartRebuild, + BaseLSN: 500, + TargetLSN: 1000, + } + if err := SendSessionControl(conn, ctrl); err != nil { + t.Fatalf("send control: %v", err) + } + + // Read ack. + msgType, payload, err := ReadFrame(conn) + if err != nil { + t.Fatalf("read ack: %v", err) + } + if msgType != MsgSessionAck { + t.Fatalf("unexpected ack type: 0x%02x", msgType) + } + ack, err := DecodeSessionAck(payload) + if err != nil { + t.Fatalf("decode ack: %v", err) + } + if ack.SessionID != 99 || ack.Phase != SessionAckAccepted { + t.Fatalf("ack mismatch: %+v", ack) + } + + if err := <-done; err != nil { + t.Fatalf("server: %v", err) + } + t.Log("session control over TCP: round-trip verified") +} + +func createTestVolForTransport(t *testing.T, name string) *BlockVol { + t.Helper() + opts := CreateOptions{ + VolumeSize: 4 * 1024 * 1024, + BlockSize: 4096, + WALSize: 1 * 1024 * 1024, + } + vol, err := CreateBlockVol(t.TempDir()+"/"+name+".blk", opts) + if err != nil { + t.Fatal(err) + } + return vol +}