mdb/kvstore/wal/writer.go

package wal

import (
	"database/sql"
	"encoding/binary"
	"log"
	"net"
	"sync"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type insertJob struct {
	Collection string
	ID         uint64
	Store      bool
	Data       []byte
	Ready      *sync.WaitGroup
}

type Writer struct {
	primary bool
	db      *sql.DB
	insert  *sql.Stmt

	lock    sync.Mutex
	running bool
	insertQ chan insertJob
	doneWG  sync.WaitGroup

	recvLock sync.Mutex
}

func NewWriterPrimary(walPath string) *Writer {
	return newWriter(walPath, true)
}

func NewWriterSecondary(walPath string) *Writer {
	return newWriter(walPath, false)
}

func newWriter(walPath string, primary bool) *Writer {
	db := initWAL(walPath)

	insert, err := db.Prepare(sqlWALInsert)
	must(err)
	w := &Writer{
		primary: primary,
		db:      db,
		insert:  insert,
	}

	if primary {
		w.start()
	}

	return w
}

func (w *Writer) Close() {
	if w.db == nil {
		return
	}

	w.stop()
	w.db.Close()
	w.db = nil
}

func (w *Writer) Store(collection string, id uint64, data []byte) {
	if !w.primary {
		panic("Store called on secondary.")
	}

	job := insertJob{
		Collection: collection,
		ID:         id,
		Store:      true,
		Data:       data,
		Ready:      &sync.WaitGroup{},
	}
	job.Ready.Add(1)
	w.insertQ <- job
	job.Ready.Wait()
}

func (w *Writer) Delete(collection string, id uint64) {
	if !w.primary {
		panic("Delete called on secondary.")
	}

	job := insertJob{
		Collection: collection,
		ID:         id,
		Store:      false,
		Ready:      &sync.WaitGroup{},
	}
	job.Ready.Add(1)
	w.insertQ <- job
	job.Ready.Wait()
}

// Called single-threaded from RecvWAL.
func (w *Writer) storeAsync(collection string, id uint64, data []byte) {
	w.insertQ <- insertJob{
		Collection: collection,
		ID:         id,
		Store:      true,
		Data:       data,
	}
}

// Called single-threaded from RecvWAL.
func (w *Writer) deleteAsync(collection string, id uint64) {
	w.insertQ <- insertJob{
		Collection: collection,
		ID:         id,
		Store:      false,
	}
}

func (w *Writer) MaxSeqNum() (n uint64) {
	w.db.QueryRow(sqlWALMaxSeqNum).Scan(&n)
	return
}

func (w *Writer) RecvWAL(conn net.Conn) {
	defer conn.Close()

	if w.primary {
		panic("RecvWAL called on primary.")
	}

	if !w.recvLock.TryLock() {
		log.Printf("Multiple calls to RecvWAL. Dropping connection.")
		return
	}
	defer w.recvLock.Unlock()

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

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

	// Send fromID to the conn.
	conn.SetWriteDeadline(time.Now().Add(connTimeout))
	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{})

	// Start processing inserts.
	w.start()
	defer w.stop()

	for {
		conn.SetReadDeadline(time.Now().Add(connTimeout))
		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)
		}
	}
}

func (w *Writer) DeleteBefore(seconds int64) {
	_, err := w.db.Exec(sqlWALDeleteQuery, time.Now().Unix()-seconds)
	must(err)
}