ntfy/message/tail.go

package message

import (
	"database/sql"
	"encoding/json"
	"sync"
	"sync/atomic"
	"time"

	"heckel.io/ntfy/v2/model"
)

const (
	tailPageSize    = 1000        // Rows per query
	tailMaxPages    = 10          // Queries per read at most, so a read drains a backlog of up to 10,000 rows
	tailHoleTimeout = time.Minute // How long a skipped row id is re-checked before it counts as never committed (rolled back)
	tailMaxHoles    = 10000       // Bounds the skipped row ids re-checked on every call
)

// Tail reads the messages written to the cache in write order (by row id), so a node can find
// what it never received live. Row ids are handed out on insert but only become visible on commit,
// and concurrent writers (cluster nodes) commit out of order: a row id the tail had to skip is a
// hole, and is re-checked until its row shows up or the hole times out.
type Tail struct {
	cache    *Cache
	cursor   int64               // Highest row id read so far
	holes    map[int64]time.Time // Skipped row ids below the cursor, and when they were skipped
	lastRead atomic.Int64        // Unix nanoseconds of the last successful read (see LastRead); 0 before the first
	done     chan struct{}       // Closed by Close, to stop the loop calling Next
	closed   sync.Once           // Guards closing done
}

// Next returns the published messages written since the last call, in write order, reading
// pages until it is caught up or has read tailMaxPages. It reads the primary: a replica may lag,
// and what it does not show yet would be skipped as a hole.
func (t *Tail) Next() ([]*model.Message, error) {
	defer t.expireHoles()
	messages := make([]*model.Message, 0)
	for i := 0; i < tailMaxPages; i++ {
		page, full, err := t.nextPage()
		if err != nil {
			if len(messages) > 0 {
				t.lastRead.Store(time.Now().UnixNano())
				return messages, nil // The pages read so far are done; the failed one is retried next time
			}
			return nil, err
		}
		messages = append(messages, page...)
		if !full {
			break
		}
	}
	t.lastRead.Store(time.Now().UnixNano())
	return messages, nil
}

// nextPage reads one page of rows: those after the cursor, plus the holes (row ids skipped
// earlier because they were not committed yet). It returns the published messages among them, and
// true if the page was full, so more may be waiting. For each row:
//   - a hole that is now visible is "filled": it is removed from the holes
//   - a row after the cursor moves the cursor to it; the ids it jumped over are "skipped", and
//     become holes
//
// Example: cursor 103, holes {104}; the page returns rows 104, 105 and 108. Row 104 fills its
// hole, 105 moves the cursor, and 108 moves it again, skipping 106 and 107. Afterwards: cursor
// 108, holes {106, 107}. Nothing changes if the read fails part way, so it is retried whole.
func (t *Tail) nextPage() ([]*model.Message, bool, error) {
	holes := make([]int64, 0, len(t.holes))
	for id := range t.holes {
		holes = append(holes, id)
	}
	holesJSON, err := json.Marshal(holes) // One parameter for any number of holes, so the statement text never changes
	if err != nil {
		return nil, false, err
	}
	rows, err := t.cache.db.Primary().Query(t.cache.queries.selectMessagesAfterRowID, t.cursor, string(holesJSON), tailPageSize)
	if err != nil {
		return nil, false, err
	}
	defer rows.Close()
	cursor, count := t.cursor, 0                          // Applied to t only once the whole page was read
	filled, skipped := make([]int64, 0), make([]int64, 0) // Holes now visible; row ids jumped over (new holes)
	messages := make([]*model.Message, 0)                 // Published messages, in row id order
	for rows.Next() {
		m, rowID, published, err := readTailRow(rows)
		if err != nil {
			return nil, false, err
		}
		count++
		if _, ok := t.holes[rowID]; ok {
			filled = append(filled, rowID)
		} else if rowID > cursor {
			// Row ids between the cursor and this one are not visible (yet)
			for id := cursor + 1; id < rowID && len(skipped) < tailMaxHoles; id++ {
				skipped = append(skipped, id)
			}
			cursor = rowID
		}
		if published { // Scheduled messages go out through the delayed sender
			messages = append(messages, m)
		}
	}
	if err := rows.Err(); err != nil {
		return nil, false, err
	}
	// Only now: a read that fails part way is retried whole, from the same position
	now := time.Now()
	for _, id := range filled {
		delete(t.holes, id)
	}
	for _, id := range skipped {
		if len(t.holes) >= tailMaxHoles {
			break
		}
		t.holes[id] = now
	}
	t.cursor = cursor
	return messages, count == tailPageSize, nil
}

// readTailRow reads a row of selectMessagesAfterRowID: the message, then its row id and whether
// it is published (false for a scheduled message)
func readTailRow(rows *sql.Rows) (*model.Message, int64, bool, error) {
	var rowID int64
	var published bool
	m, err := readMessage(rows, &rowID, &published)
	if err != nil {
		return nil, 0, false, err
	}
	return m, rowID, published, nil
}

func (t *Tail) expireHoles() {
	for id, skipped := range t.holes {
		if time.Since(skipped) > tailHoleTimeout {
			delete(t.holes, id)
		}
	}
}

// LastRead returns when the tail last read the database successfully; zero before the first
// read. Unlike Next, it may be called from any goroutine.
func (t *Tail) LastRead() time.Time {
	nanos := t.lastRead.Load()
	if nanos == 0 {
		return time.Time{}
	}
	return time.Unix(0, nanos)
}

// Done returns a channel that is closed when the tail is closed, so the loop calling Next can stop
func (t *Tail) Done() <-chan struct{} {
	return t.done
}

// Close tells the loop calling Next to stop (see Done); it may be called more than once
func (t *Tail) Close() {
	t.closed.Do(func() {
		close(t.done)
	})
}

// addHoles records the row ids between from and to (both exclusive) as holes
func (t *Tail) addHoles(from, to int64, now time.Time) {
	for id := from + 1; id < to && len(t.holes) < tailMaxHoles; id++ {
		t.holes[id] = now
	}
}