pizzasql/pkg/executor/notify.go

package executor

import (
	"fmt"

	"github.com/danfragoso/pizzasql-next/pkg/parser"
	"github.com/danfragoso/pizzasql-next/pkg/storage"
)

// listenOp records one LISTEN/UNLISTEN statement so its effect can be applied
// at transaction commit. PostgreSQL makes subscription changes visible only
// after commit; a rollback leaves the previous subscription set untouched.
type listenOp struct {
	kind    listenOpKind
	channel string
}

type listenOpKind int

const (
	listenOpListen listenOpKind = iota
	listenOpUnlisten
	listenOpUnlistenAll
)

// notificationSnapshot captures the staged notification state at a savepoint.
type notificationSnapshot struct {
	listenOps   int
	notifyQueue int
}

// SetNotificationHub installs the per-database notification registry and
// creates this session's subscriber. It must be called before the executor
// processes LISTEN or NOTIFY. Safe to call once per connection.
func (e *Executor) SetNotificationHub(hub *storage.Hub) {
	e.hub = hub
	if hub == nil {
		return
	}
	if e.subscriber == nil {
		e.subscriber = storage.NewSubscriber(e.backendPID, 0)
	}
}

// SetBackendPID sets the backend PID reported in NotificationResponse frames.
// It also applies to an already-created subscriber.
func (e *Executor) SetBackendPID(pid int32) {
	e.backendPID = pid
	if e.subscriber == nil && e.hub != nil {
		e.subscriber = storage.NewSubscriber(pid, 0)
	}
}

// NotificationChannel returns the receive side of this session's notification
// queue, or nil when no hub is configured.
func (e *Executor) NotificationChannel() <-chan storage.Notification {
	if e.subscriber == nil {
		return nil
	}
	return e.subscriber.Notifications()
}

// NotificationOverflow is closed when the session's outbound queue is full and
// the hub has dropped it. Protocol layers should disconnect the connection.
func (e *Executor) NotificationOverflow() <-chan struct{} {
	if e.subscriber == nil {
		return nil
	}
	return e.subscriber.Overflow()
}

// Close releases this session's subscriptions. It is the disconnect path; it
// is safe to call more than once.
func (e *Executor) Close() error {
	if e.hub != nil && e.subscriber != nil {
		e.hub.UnsubscribeAll(e.subscriber)
	}
	e.listens = make(map[string]bool)
	e.listenOps = nil
	e.notifyQueue = nil
	e.notificationSaves = nil
	return e.RollbackActive()
}

func (e *Executor) executeListen(stmt *parser.ListenStmt) (*Result, error) {
	if stmt.Channel == "" {
		return nil, fmt.Errorf("channel name cannot be empty")
	}
	if e.inTransaction {
		e.listenOps = append(e.listenOps, listenOp{kind: listenOpListen, channel: stmt.Channel})
		return NewResult("LISTEN"), nil
	}
	e.applyListenOp(listenOp{kind: listenOpListen, channel: stmt.Channel})
	return NewResult("LISTEN"), nil
}

func (e *Executor) executeUnlisten(stmt *parser.UnlistenStmt) (*Result, error) {
	op := listenOp{kind: listenOpUnlisten, channel: stmt.Channel}
	if stmt.All {
		op = listenOp{kind: listenOpUnlistenAll}
	} else if stmt.Channel == "" {
		return nil, fmt.Errorf("channel name cannot be empty")
	}
	if e.inTransaction {
		e.listenOps = append(e.listenOps, op)
		return NewResult("UNLISTEN"), nil
	}
	e.applyListenOp(op)
	return NewResult("UNLISTEN"), nil
}

func (e *Executor) executeNotify(stmt *parser.NotifyStmt) (*Result, error) {
	if len(stmt.Payload) > 8000 {
		return nil, fmt.Errorf("notification payload is longer than 8000 bytes")
	}
	n := storage.Notification{
		PID:     e.backendPID,
		Channel: stmt.Channel,
		Payload: stmt.Payload,
	}
	if !e.inTransaction {
		e.publish(n)
		return NewResult("NOTIFY"), nil
	}
	// PostgreSQL collapses duplicate (channel, payload) pairs within one
	// transaction, regardless of how many times they were raised.
	for _, queued := range e.notifyQueue {
		if queued.Channel == n.Channel && queued.Payload == n.Payload {
			return NewResult("NOTIFY"), nil
		}
	}
	e.notifyQueue = append(e.notifyQueue, n)
	return NewResult("NOTIFY"), nil
}

func (e *Executor) applyListenOp(op listenOp) {
	if e.hub == nil || e.subscriber == nil {
		return
	}
	switch op.kind {
	case listenOpListen:
		if !e.listens[op.channel] {
			e.hub.Subscribe(e.subscriber, op.channel)
			e.listens[op.channel] = true
		}
	case listenOpUnlisten:
		if e.listens[op.channel] {
			e.hub.Unsubscribe(e.subscriber, op.channel)
			delete(e.listens, op.channel)
		}
	case listenOpUnlistenAll:
		for channel := range e.listens {
			e.hub.Unsubscribe(e.subscriber, channel)
		}
		e.listens = make(map[string]bool)
	}
}

func (e *Executor) publish(n storage.Notification) {
	if e.hub == nil {
		return
	}
	e.hub.Publish(n)
}

// finishCommit applies staged subscription changes and releases queued
// notifications. It runs only after the storage commit has succeeded.
func (e *Executor) finishCommit(lsn uint64) {
	for _, op := range e.listenOps {
		e.applyListenOp(op)
	}
	for _, n := range e.notifyQueue {
		n.LSN = lsn
		e.publish(n)
	}
	e.clearNotificationState()
}

// finishRollback discards staged subscription changes and notifications.
func (e *Executor) finishRollback() {
	e.clearNotificationState()
}

func (e *Executor) clearNotificationState() {
	e.listenOps = nil
	e.notifyQueue = nil
	e.notificationSaves = nil
}

func (e *Executor) notificationSnapshot() notificationSnapshot {
	return notificationSnapshot{
		listenOps:   len(e.listenOps),
		notifyQueue: len(e.notifyQueue),
	}
}

func (e *Executor) rollbackNotificationsTo(s notificationSnapshot) {
	if s.listenOps <= len(e.listenOps) {
		e.listenOps = e.listenOps[:s.listenOps]
	}
	if s.notifyQueue <= len(e.notifyQueue) {
		e.notifyQueue = e.notifyQueue[:s.notifyQueue]
	}
}