consumer.go

package wal

import (
	"context"
	"errors"
	"time"

	"github.com/mgurevin/wal/internal/lockrank"
)

type (
	// Message is one borrowed WAL record offered to a downstream.
	// Senders must not retain or modify Data after Send returns.
	Message struct {
		Position Position
		ID       string
		Data     []byte
	}

	// Batch is one bounded contiguous WAL prefix. Senders must not retain or
	// modify Messages or their Data after Send returns.
	Batch struct {
		Messages []Message
	}

	// DeliveryGuarantee selects one complete source and acknowledgment profile.
	DeliveryGuarantee uint8

	// PermanentRejectionPolicy selects an explicit rejected-record disposition.
	PermanentRejectionPolicy uint8

	// Rejection identifies one permanently rejected message. Cause and ID may
	// contain private application information.
	Rejection struct {
		ID    string
		Cause error
	}

	// SendResult describes one durably accepted prefix and an optional
	// immediately following permanent rejection.
	SendResult struct {
		ThroughID string
		Rejection *Rejection
	}

	// Sender delivers one batch and returns its durable downstream result.
	Sender interface {
		Send(context.Context, Batch) (SendResult, error)
	}

	// SenderFunc adapts a function to Sender.
	SenderFunc func(context.Context, Batch) (SendResult, error)

	// RejectionHandler durably and idempotently accepts one rejected message.
	RejectionHandler interface {
		Handle(context.Context, Message, Rejection) error
	}

	// RejectionHandlerFunc adapts a function to RejectionHandler.
	RejectionHandlerFunc func(context.Context, Message, Rejection) error

	// Acknowledger durably commits one accepted WAL prefix. Implementations may
	// coordinate additional application state before truncating the log. After
	// downstream acceptance, Consumer preserves context values but removes
	// cancellation until Acknowledge returns a definite outcome.
	Acknowledger interface {
		Supports(Durability) bool
		Acknowledge(context.Context, Position, Durability) error
	}

	// ConsumerStats is one lock-consistent delivery progress snapshot.
	ConsumerStats struct {
		Head                Position
		HeadAttempts        uint64
		HeadSince           time.Time
		LastProgress        time.Time
		Stalled             bool
		PermanentRejections uint64
		DeadLettered        uint64
		Discarded           uint64
	}

	// ConsumerConfig defines one bounded FIFO consumer.
	ConsumerConfig struct {
		// Log is the ordered source. The consumer never closes it.
		Log Log
		// Acknowledger commits local acknowledgments. Nil truncates Log directly.
		Acknowledger Acknowledger
		// Sender owns the downstream transport and durable acceptance contract.
		Sender Sender
		// Guarantee selects the source and acknowledgment durability profile.
		// Zero defaults to CrashPersistent.
		Guarantee DeliveryGuarantee
		// RejectionPolicy selects permanent rejection handling.
		RejectionPolicy PermanentRejectionPolicy
		// RejectionHandler is required only for dead-letter handling.
		RejectionHandler RejectionHandler
		// MaxBatch bounds message headers offered in one call. Zero defaults
		// to 128 and cannot exceed MaxConsumerBatchEntries.
		MaxBatch int
		// MaxBatchBytes bounds combined opaque payload bytes in one call. Zero
		// defaults to 4 MiB.
		MaxBatchBytes int
		// IdleDelay controls empty-WAL polling. Zero defaults to 250 ms.
		IdleDelay time.Duration
		// RetryDelay is the first downstream failure delay. Zero defaults to
		// one second.
		RetryDelay time.Duration
		// MaxRetryDelay caps exponential downstream retry delay. Zero defaults
		// to 30 seconds.
		MaxRetryDelay time.Duration
		// StallAfter marks an unchanged head stalled after this duration. Zero
		// defaults to one minute.
		StallAfter time.Duration
		// StallAfterAttempts marks an unchanged head stalled after this many
		// failed attempts. Zero defaults to 10.
		StallAfterAttempts uint64
		// OnInternalError observes contained callbacks and retryable downstream
		// failures. A panic in the reporter is contained.
		OnInternalError func(error)
		// Logf is the fallback when OnInternalError is absent or panics.
		// Nil or a panic in Logf falls back to log.Printf.
		Logf func(string, ...any)
	}

	// Consumer delivers at most one bounded borrowed WAL batch at a time.
	// Run, Drain, and Next are mutually exclusive.
	Consumer struct {
		// operation owns one complete delivery attempt, including caller
		// callbacks. statsMutex is nested only for short counter updates.
		operation        lockrank.Mutex
		log              Log
		acknowledger     Acknowledger
		sender           Sender
		rejectionHandler RejectionHandler
		durability       Durability
		rejectionPolicy  PermanentRejectionPolicy
		maxBatch         int
		maxBatchBytes    int
		idleDelay        time.Duration
		retryDelay       time.Duration
		maxRetryDelay    time.Duration
		stallAfter       time.Duration
		stallAttempts    uint64
		onInternalError  func(error)
		logf             func(string, ...any)
		now              func() time.Time
		statsMutex       lockrank.Mutex
		stats            ConsumerStats
		stallReported    bool
		messages         []Message
		payloads         []byte
		payloadEnds      []int
	}

	sendError struct {
		operation string
		cause     error
	}

	// PermanentRejectionError contains private rejected-record diagnostics.
	PermanentRejectionError struct {
		rejection Rejection
	}

	operationError struct {
		message string
		cause   error
	}

	logAcknowledger struct {
		log Log
	}
)

const (
	// ProcessLifetime uses volatile source and acknowledgment durability.
	ProcessLifetime DeliveryGuarantee = iota + 1
	// CrashPersistent requires synchronized source and acknowledgment support.
	CrashPersistent
)

const (
	// StopOnPermanentRejection retains the rejected head and stops delivery.
	StopOnPermanentRejection PermanentRejectionPolicy = iota
	// DeadLetterPermanentRejection invokes the configured durable handler.
	DeadLetterPermanentRejection
	// DiscardPermanentRejection explicitly truncates the rejected record.
	DiscardPermanentRejection
)

const (
	defaultConsumerIdleDelay     = 250 * time.Millisecond
	defaultConsumerRetryDelay    = time.Second
	defaultConsumerMaxRetryDelay = 30 * time.Second
	defaultConsumerMaxBatch      = 128
	defaultConsumerMaxBatchBytes = 4 << 20
	defaultConsumerStallAfter    = time.Minute
	defaultConsumerStallAttempts = 10
	// MaxConsumerBatchEntries is the hard allocation-safety ceiling.
	MaxConsumerBatchEntries = 4_096
)

var (
	// ErrUnsupportedDeliveryGuarantee reports a profile unsupported by either
	// the source or acknowledgment committer.
	ErrUnsupportedDeliveryGuarantee = errors.New(
		"wal consumer: unsupported delivery guarantee; choose ProcessLifetime or use Synced-capable dependencies",
	)
	// ErrBatchTooLarge reports a WAL record larger than MaxBatchBytes.
	ErrBatchTooLarge = errors.New(
		"wal consumer: one WAL record exceeds MaxBatchBytes; raise the bound or store a smaller record",
	)
	// ErrInvalidSendResult reports a non-contiguous acceptance or rejection.
	ErrInvalidSendResult = errors.New(
		"wal consumer: invalid downstream result; return one contiguous accepted prefix and optional next rejection",
	)
	// ErrPermanentRejection reports a retained permanently rejected head.
	ErrPermanentRejection = errors.New(
		"wal consumer: permanent rejection retained; resolve it or configure a reviewed dead-letter policy",
	)
	// ErrConsumerStalled reports one unchanged-head observability transition.
	ErrConsumerStalled = errors.New(
		"wal consumer: delivery head stalled; inspect the downstream and retained head before retrying",
	)
	// ErrRunning reports concurrent use of Next, Drain, or Run.
	ErrRunning = errors.New(
		"wal consumer: another delivery operation is running; serialize Next, Drain, and Run per instance",
	)
	errBatchFull         = errors.New("wal consumer: internal batch byte limit reached")
	errConsumerWatermark = errors.New("wal consumer: internal durable watermark reached")
)

// NewConsumer constructs a consumer. It starts no goroutine.
func NewConsumer(config ConsumerConfig) (*Consumer, error) {
	if config.Log == nil {
		return nil, errors.New(
			"wal consumer: Log is required; configure the WAL to deliver",
		)
	}

	if config.Sender == nil {
		return nil, errors.New(
			"wal consumer: Sender is required; configure a durable downstream sender",
		)
	}

	guarantee := config.Guarantee
	if guarantee == 0 {
		guarantee = CrashPersistent
	}

	var durability Durability

	switch guarantee {
	case ProcessLifetime:
		durability = Volatile

	case CrashPersistent:
		durability = Synced

	default:
		return nil, ErrUnsupportedDeliveryGuarantee
	}

	acknowledger := config.Acknowledger
	if acknowledger == nil {
		acknowledger = logAcknowledger{log: config.Log}
	}

	if !config.Log.Supports(durability) ||
		!acknowledger.Supports(durability) {
		return nil, ErrUnsupportedDeliveryGuarantee
	}

	switch config.RejectionPolicy {
	case StopOnPermanentRejection,
		DiscardPermanentRejection:
		if config.RejectionHandler != nil {
			return nil, errors.New(
				"wal consumer: rejection handler requires dead-letter policy; remove it or select DeadLetterPermanentRejection",
			)
		}

	case DeadLetterPermanentRejection:
		if config.RejectionHandler == nil {
			return nil, errors.New(
				"wal consumer: dead-letter policy requires rejection handler; configure a durable idempotent RejectionHandler",
			)
		}

	default:
		return nil, errors.New(
			"wal consumer: invalid permanent rejection policy; choose stop, dead-letter, or explicit discard",
		)
	}

	maxBatch := config.MaxBatch
	if maxBatch == 0 {
		maxBatch = defaultConsumerMaxBatch
	}

	maxBatchBytes := config.MaxBatchBytes
	if maxBatchBytes == 0 {
		maxBatchBytes = defaultConsumerMaxBatchBytes
	}

	if maxBatch < 0 ||
		maxBatch > MaxConsumerBatchEntries ||
		maxBatchBytes < 0 {
		return nil, errors.New(
			"wal consumer: batch limits must not be negative and MaxBatch must not exceed MaxConsumerBatchEntries; use zero for defaults",
		)
	}

	idleDelay := defaultDuration(config.IdleDelay, defaultConsumerIdleDelay)
	retryDelay := defaultDuration(config.RetryDelay, defaultConsumerRetryDelay)
	maxRetryDelay := defaultDuration(
		config.MaxRetryDelay,
		defaultConsumerMaxRetryDelay,
	)
	stallAfter := defaultDuration(
		config.StallAfter,
		defaultConsumerStallAfter,
	)

	stallAttempts := config.StallAfterAttempts
	if stallAttempts == 0 {
		stallAttempts = defaultConsumerStallAttempts
	}

	if idleDelay < 0 ||
		retryDelay < 0 ||
		maxRetryDelay < retryDelay ||
		stallAfter < 0 {
		return nil, errors.New(
			"wal consumer: delays must be positive and ordered; use zero for defaults and keep MaxRetryDelay no shorter than RetryDelay",
		)
	}

	return &Consumer{
		operation:        lockrank.NewMutex(lockrank.FIFOOperation),
		statsMutex:       lockrank.NewMutex(lockrank.FIFOStats),
		log:              config.Log,
		acknowledger:     acknowledger,
		sender:           config.Sender,
		rejectionHandler: config.RejectionHandler,
		durability:       durability,
		rejectionPolicy:  config.RejectionPolicy,
		maxBatch:         maxBatch,
		maxBatchBytes:    maxBatchBytes,
		idleDelay:        idleDelay,
		retryDelay:       retryDelay,
		maxRetryDelay:    maxRetryDelay,
		stallAfter:       stallAfter,
		stallAttempts:    stallAttempts,
		onInternalError:  config.OnInternalError,
		logf:             config.Logf,
		now:              time.Now,
		messages:         make([]Message, 0, min(maxBatch, defaultConsumerMaxBatch)),
		payloadEnds:      make([]int, 0, min(maxBatch, defaultConsumerMaxBatch)),
	}, nil
}

// Send implements SenderFunc.
func (f SenderFunc) Send(
	ctx context.Context,
	batch Batch,
) (SendResult, error) {
	return f(ctx, batch)
}

// Handle implements RejectionHandler.
func (f RejectionHandlerFunc) Handle(
	ctx context.Context,
	message Message,
	rejection Rejection,
) error {
	return f(ctx, message, rejection)
}

// Next offers one bounded live WAL batch and returns the number removed from
// the local head after primary acceptance or the configured rejection
// disposition. It reports zero without error when the durable source
// watermark is empty.
func (c *Consumer) Next(ctx context.Context) (int, error) {
	if !c.operation.TryLock() {
		return 0, ErrRunning
	}
	defer c.operation.Unlock()

	return c.next(ctx)
}

// Drain advances records until the durable source watermark is empty or an
// operation fails. It does not retry downstream failures.
func (c *Consumer) Drain(ctx context.Context) (int, error) {
	if !c.operation.TryLock() {
		return 0, ErrRunning
	}
	defer c.operation.Unlock()

	advanced := 0

	for {
		count, err := c.next(ctx)
		advanced += count

		if err != nil || count == 0 {
			return advanced, err
		}
	}
}

// Run polls an empty WAL and retries downstream errors with bounded
// exponential backoff until ctx ends. Local persistence, protocol, and
// callback-panic errors stop the worker so the caller can inspect or recover
// local state.
func (c *Consumer) Run(ctx context.Context) error {
	if !c.operation.TryLock() {
		return ErrRunning
	}
	defer c.operation.Unlock()

	retryDelay := c.retryDelay

	for {
		advanced, err := c.next(ctx)
		if err == nil {
			retryDelay = c.retryDelay

			if advanced != 0 {
				continue
			}

			if err := wait(ctx, c.idleDelay); err != nil {
				return err
			}

			continue
		}

		var senderFailure *sendError
		if !errors.As(err, &senderFailure) {
			return err
		}

		if advanced != 0 {
			retryDelay = c.retryDelay
		}

		c.report(err)

		if err := wait(ctx, retryDelay); err != nil {
			return err
		}

		retryDelay = nextDelay(retryDelay, c.maxRetryDelay)
	}
}

func (c *Consumer) next(ctx context.Context) (int, error) {
	if err := ctx.Err(); err != nil {
		return 0, err
	}

	watermark, err := c.log.Flush(ctx, c.durability)
	if err != nil {
		return 0, &operationError{
			message: "wal consumer: cannot establish the source durability watermark; inspect local storage",
			cause:   err,
		}
	}

	state := c.log.State()
	c.observeHead(state, c.now())

	if state.Truncated >= watermark {
		return 0, nil
	}

	c.messages = c.messages[:0]
	c.payloads = c.payloads[:0]

	c.payloadEnds = c.payloadEnds[:0]
	defer func() {
		clear(c.messages)
		c.messages = c.messages[:0]
		clear(c.payloads)
		c.payloads = c.payloads[:0]
		clear(c.payloadEnds)
		c.payloadEnds = c.payloadEnds[:0]
	}()

	batchBytes := 0

	err = c.log.Scan(
		ctx,
		state.Truncated,
		c.maxBatch,
		func(entry Entry) error {
			if entry.Position > watermark {
				return errConsumerWatermark
			}

			if len(entry.Record.Data) > c.maxBatchBytes {
				return ErrBatchTooLarge
			}

			if len(c.messages) != 0 &&
				len(entry.Record.Data) > c.maxBatchBytes-batchBytes {
				return errBatchFull
			}

			c.messages = append(c.messages, Message{
				Position: entry.Position,
				ID:       entry.Record.ID,
			})
			// Scan borrows one record at a time, while Sender owns one stable
			// batch for the complete synchronous call. Copying here preserves
			// that lifetime without coupling Consumer to the file decoder.
			// Native qualification found no Synced-producer gain from the
			// streaming alternative; DESIGN.md records the measured decision.
			c.payloads = append(c.payloads, entry.Record.Data...)
			c.payloadEnds = append(c.payloadEnds, len(c.payloads))
			batchBytes += len(entry.Record.Data)

			return nil
		},
	)
	if err != nil &&
		!errors.Is(err, errBatchFull) &&
		!errors.Is(err, errConsumerWatermark) {
		return 0, &operationError{
			message: "wal consumer: cannot read the next WAL batch; inspect local storage",
			cause:   err,
		}
	}

	if len(c.messages) == 0 {
		return 0, errors.New("wal consumer: WAL state changed without a readable record; retry after storage inspection")
	}

	payloadOffset := 0

	for index := range c.messages {
		payloadEnd := c.payloadEnds[index]
		c.messages[index].Data = c.payloads[payloadOffset:payloadEnd]
		payloadOffset = payloadEnd
	}

	result, err := send(c.sender, ctx, Batch{Messages: c.messages})
	if err != nil {
		c.recordHeadFailure(c.now())

		return 0, err
	}

	resolution, err := validateSendResult(c.messages, result)
	if err != nil {
		c.recordHeadFailure(c.now())

		return 0, err
	}

	advanced := 0

	if resolution.accepted != 0 {
		through := c.messages[resolution.accepted-1].Position
		if err := c.acknowledger.Acknowledge(
			withoutCancellation(ctx),
			through,
			c.durability,
		); err != nil {
			c.recordHeadFailure(c.now())

			return 0, &operationError{
				message: "wal consumer: downstream accepted a prefix but local truncation failed; retry after storage recovery",
				cause:   err,
			}
		}

		advanced = resolution.accepted

		c.recordProgress(c.log.State(), c.now())
	}

	if resolution.rejected < 0 {
		return advanced, nil
	}

	message := c.messages[resolution.rejected]
	rejection := *result.Rejection

	c.recordPermanentRejection()

	switch c.rejectionPolicy {
	case StopOnPermanentRejection:
		c.recordHeadFailure(c.now())
		c.recordPermanentStop()

		return advanced, newPermanentRejectionError(rejection)

	case DeadLetterPermanentRejection:
		if err := invokeRejectionHandler(
			c.rejectionHandler,
			ctx,
			message,
			rejection,
		); err != nil {
			c.recordHeadFailure(c.now())

			if isCallbackPanic(err) {
				return advanced, err
			}

			return advanced, &sendError{
				operation: "permanent rejection handler",
				cause:     err,
			}
		}

		if err := c.truncateRejected(ctx, message.Position); err != nil {
			return advanced, err
		}

		c.recordDeadLetter()

	case DiscardPermanentRejection:
		if err := c.truncateRejected(ctx, message.Position); err != nil {
			return advanced, err
		}

		c.recordDiscard()
	}

	c.recordProgress(c.log.State(), c.now())

	return advanced + 1, nil
}

func (c *Consumer) truncateRejected(
	ctx context.Context,
	position Position,
) error {
	if err := c.acknowledger.Acknowledge(
		withoutCancellation(ctx),
		position,
		c.durability,
	); err != nil {
		c.recordHeadFailure(c.now())

		return &operationError{
			message: "wal consumer: rejected record disposition succeeded but local truncation failed; retry after storage recovery",
			cause:   err,
		}
	}

	return nil
}

func (c *Consumer) report(err error) {
	reportInternalError(c.onInternalError, c.logf, err)
}

func (a logAcknowledger) Supports(durability Durability) bool {
	return a.log.Supports(durability)
}

func (a logAcknowledger) Acknowledge(
	ctx context.Context,
	through Position,
	durability Durability,
) error {
	return a.log.Truncate(ctx, through, durability)
}

func (e *sendError) Error() string {
	if e.operation != "" {
		return "wal consumer: " + e.operation +
			" failed; retry with the same message ID"
	}

	return "wal consumer: downstream send failed; retry with the same message IDs"
}

func (e *sendError) Unwrap() error {
	return e.cause
}

func (e *operationError) Error() string {
	return e.message
}

func (e *operationError) Unwrap() error {
	return e.cause
}

func send(
	sender Sender,
	ctx context.Context,
	batch Batch,
) (result SendResult, err error) {
	defer func() {
		if recovered := recover(); recovered != nil {
			result = SendResult{}
			err = newCallbackPanicError(
				"downstream sender",
				recovered,
			)
		}
	}()

	result, err = sender.Send(ctx, batch)
	if err != nil {
		return SendResult{}, &sendError{cause: err}
	}

	return result, nil
}

func defaultDuration(value, fallback time.Duration) time.Duration {
	if value == 0 {
		return fallback
	}

	return value
}

func nextDelay(current, maximum time.Duration) time.Duration {
	if current >= maximum || current > maximum/2 {
		return maximum
	}

	return current * 2
}

func wait(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()

	select {
	case <-ctx.Done():
		return ctx.Err()

	case <-timer.C:
		return nil
	}
}