mdb/wal/shipping.go

package wal

import (
	"encoding/binary"
	"log"
	"net"
	"time"
)

const recHeaderSize = 22

func encodeRecordHeader(rec Record, buf []byte) {
	// SeqNum (8)
	// ID (8)
	// DataLen (4)
	// Store (1)
	// CollectionLen (1)

	binary.LittleEndian.PutUint64(buf[:8], rec.SeqNum)
	buf = buf[8:]
	binary.LittleEndian.PutUint64(buf[:8], rec.ID)
	buf = buf[8:]

	if rec.Store {
		binary.LittleEndian.PutUint32(buf[:4], uint32(len(rec.Data)))
		buf[4] = 1
	} else {
		binary.LittleEndian.PutUint32(buf[:4], 0)
		buf[4] = 0
	}
	buf = buf[5:]

	buf[0] = byte(len(rec.Collection))
}

func decodeRecHeader(header []byte) (rec Record, colLen, dataLen int) {
	buf := header

	rec.SeqNum = binary.LittleEndian.Uint64(buf[:8])
	buf = buf[8:]
	rec.ID = binary.LittleEndian.Uint64(buf[:8])
	buf = buf[8:]
	dataLen = int(binary.LittleEndian.Uint32(buf[:4]))
	buf = buf[4:]
	rec.Store = buf[0] != 0
	buf = buf[1:]
	colLen = int(buf[0])

	return
}

func SendWAL(walPath string, conn net.Conn) {
	defer conn.Close()

	buf := make([]byte, 8)
	headerBuf := make([]byte, recHeaderSize)
	empty := make([]byte, recHeaderSize)

	// Read the fromID from the conn.
	conn.SetReadDeadline(time.Now().Add(16 * time.Second))
	if _, err := conn.Read(buf[:8]); err != nil {
		log.Printf("SendWAL failed to read from ID: %v", err)
		return
	}

	afterSeqNum := binary.LittleEndian.Uint64(buf[:8])

	follower := NewFollower(walPath)
	defer follower.Close()

	for {
		conn.SetWriteDeadline(time.Now().Add(16 * time.Second))

		// Nothing to do.
		if follower.MaxSeqNum() <= afterSeqNum {
			if _, err := conn.Write(empty); err != nil {
				log.Printf("SendWAL failed to send heartbeat: %v", err)
				return
			}
			time.Sleep(time.Second)
			continue
		}

		err := follower.Replay(afterSeqNum, func(rec Record) error {
			afterSeqNum = rec.SeqNum
			encodeRecordHeader(rec, headerBuf)
			if _, err := conn.Write(headerBuf); err != nil {
				log.Printf("SendWAL failed to send header %v", err)
				return err
			}

			if _, err := conn.Write([]byte(rec.Collection)); err != nil {
				log.Printf("SendWAL failed to send collection name %v", err)
				return err
			}

			if !rec.Store {
				return nil
			}

			if _, err := conn.Write(rec.Data); err != nil {
				log.Printf("SendWAL failed to send data %v", err)
				return err
			}
			return nil
		})

		if err != nil {
			return
		}
	}
}

func RecvWAL(walPath string, conn net.Conn) {
	defer conn.Close()

	headerBuf := make([]byte, recHeaderSize)
	buf := make([]byte, 8)

	w := NewWriter(walPath, true)
	defer w.Close()

	afterSeqNum := w.MaxSeqNum()
	expectedSeqNum := afterSeqNum + 1

	// Send fromID to the conn.
	conn.SetWriteDeadline(time.Now().Add(time.Minute))
	binary.LittleEndian.PutUint64(buf, afterSeqNum)
	if _, err := conn.Write(buf); err != nil {
		log.Printf("RecvWAL failed to send after sequence number: %v", err)
		return
	}
	conn.SetWriteDeadline(time.Time{})

	for {
		conn.SetReadDeadline(time.Now().Add(time.Minute))
		if _, err := conn.Read(headerBuf); err != nil {
			log.Printf("RecvWAL failed to read header: %v", err)
			return
		}
		rec, colLen, dataLen := decodeRecHeader(headerBuf)

		// Heartbeat.
		if rec.SeqNum == 0 {
			continue
		}

		if rec.SeqNum != expectedSeqNum {
			log.Printf("Expected sequence number %d but got %d.",
				expectedSeqNum, rec.SeqNum)
			return
		}
		expectedSeqNum++

		if cap(buf) < colLen {
			buf = make([]byte, colLen)
		}
		buf = buf[:colLen]

		if _, err := conn.Read(buf); err != nil {
			log.Printf("RecvWAL failed to collection name: %v", err)
			return
		}

		rec.Collection = string(buf)

		if rec.Store {
			rec.Data = make([]byte, dataLen)
			if _, err := conn.Read(rec.Data); err != nil {
				log.Printf("RecvWAL failed to data: %v", err)
				return
			}
		}
		if rec.Store {
			w.storeAsync(rec.Collection, rec.ID, rec.Data)
		} else {
			w.deleteAsync(rec.Collection, rec.ID)
		}
	}
}