package pubsub

import (

"context"

"database/sql"

"database/sql/driver"

"errors"

"io"

"net"

"sync"

"sync/atomic"

"time"

"github.com/lib/pq"

"github.com/prometheus/client_golang/prometheus"

"golang.org/x/xerrors"

"cdr.dev/slog/v3"

"github.com/coder/coder/v2/coderd/database"

)

// Listener represents a pubsub handler.

type Listener func(ctx context.Context, message []byte)

// ListenerWithErr represents a pubsub handler that can also receive error

// indications

type ListenerWithErr func(ctx context.Context, message []byte, err error)

// ErrDroppedMessages is sent to ListenerWithErr if messages are dropped or

// might have been dropped.

var ErrDroppedMessages = xerrors.New("dropped messages")

// LatencyMeasureTimeout defines how often to trigger a new background latency measurement.

const LatencyMeasureTimeout = time.Second * 10

type Subscriber interface {

Subscribe(event string, listener Listener) (cancel func(), err error)

SubscribeWithErr(event string, listener ListenerWithErr) (cancel func(), err error)

}

type Publisher interface {

Publish(event string, message []byte) error

}

// Pubsub is a generic interface for broadcasting and receiving messages.

// Implementors should assume high-availability with the backing implementation.

type Pubsub interface {

Subscriber

Publisher

Close() error

}

// msgOrErr either contains a message or an error

type msgOrErr struct {

msg []byte

err error

}

// msgQueue implements a fixed length queue with the ability to replace elements

// after they are queued (but before they are dequeued).

//

// The purpose of this data structure is to build something that works a bit

// like a golang channel, but if the queue is full, then we can replace the

// last element with an error so that the subscriber can get notified that some

// messages were dropped, all without blocking.

type msgQueue struct {

ctx context.Context

cond *sync.Cond

q [BufferSize]msgOrErr

front int

size int

closed bool

l Listener

le ListenerWithErr

}

func newMsgQueue(ctx context.Context, l Listener, le ListenerWithErr) *msgQueue {

if l == nil && le == nil {

panic("l or le must be non-nil")

}

q := &msgQueue{

ctx: ctx,

cond: sync.NewCond(&sync.Mutex{}),

l: l,

le: le,

}

go q.run()

return q

}

func (q *msgQueue) run() {

for {

// wait until there is something on the queue or we are closed

q.cond.L.Lock()

for q.size == 0 && !q.closed {

q.cond.Wait()

}

if q.closed {

q.cond.L.Unlock()

return

}

item := q.q[q.front]

q.front = (q.front + 1) % BufferSize

q.size--

q.cond.L.Unlock()

// process item without holding lock

if item.err == nil {

// real message

if q.l != nil {

q.l(q.ctx, item.msg)

continue

}

if q.le != nil {

q.le(q.ctx, item.msg, nil)

continue

}

// unhittable

continue

}

// if the listener wants errors, send it.

if q.le != nil {

q.le(q.ctx, nil, item.err)

}

}

}

func (q *msgQueue) enqueue(msg []byte) {

q.cond.L.Lock()

defer q.cond.L.Unlock()

if q.size == BufferSize {

// queue is full, so we're going to drop the msg we got called with.

// We also need to record that messages are being dropped, which we

// do at the last message in the queue. This potentially makes us

// lose 2 messages instead of one, but it's more important at this

// point to warn the subscriber that they're losing messages so they

// can do something about it.

back := (q.front + BufferSize - 1) % BufferSize

q.q[back].msg = nil

q.q[back].err = ErrDroppedMessages

return

}

// queue is not full, insert the message

next := (q.front + q.size) % BufferSize

q.q[next].msg = msg

q.q[next].err = nil

q.size++

q.cond.Broadcast()

}

func (q *msgQueue) close() {

q.cond.L.Lock()

defer q.cond.L.Unlock()

defer q.cond.Broadcast()

q.closed = true

}

// dropped records an error in the queue that messages might have been dropped

func (q *msgQueue) dropped() {

q.cond.L.Lock()

defer q.cond.L.Unlock()

if q.size == BufferSize {

// queue is full, but we need to record that messages are being dropped,

// which we do at the last message in the queue. This potentially drops

// another message, but it's more important for the subscriber to know.

back := (q.front + BufferSize - 1) % BufferSize

q.q[back].msg = nil

q.q[back].err = ErrDroppedMessages

return

}

// queue is not full, insert the error

next := (q.front + q.size) % BufferSize

q.q[next].msg = nil

q.q[next].err = ErrDroppedMessages

q.size++

q.cond.Broadcast()

}

// pqListener is an interface that represents a *pq.Listener for testing

type pqListener interface {

io.Closer

Listen(string) error

Unlisten(string) error

NotifyChan() <-chan *pq.Notification

}

type pqListenerShim struct {

*pq.Listener

}

func (l pqListenerShim) NotifyChan() <-chan *pq.Notification {

return l.Notify

}

type queueSet struct {

m map[*msgQueue]struct{}

// unlistenInProgress will be non-nil if another goroutine is unlistening for the event this

// queueSet corresponds to. If non-nil, that goroutine will close the channel when it is done.

unlistenInProgress chan struct{}

}

func newQueueSet() *queueSet {

return &queueSet{

m: make(map[*msgQueue]struct{}),

}

}

// PGPubsub is a pubsub implementation using PostgreSQL.

type PGPubsub struct {

logger slog.Logger

listenDone chan struct{}

pgListener pqListener

db *sql.DB

qMu sync.Mutex

queues map[string]*queueSet

// making the close state its own mutex domain simplifies closing logic so

// that we don't have to hold the qMu --- which could block processing

// notifications while the pqListener is closing.

closeMu sync.Mutex

closedListener bool

closeListenerErr error

publishesTotal *prometheus.CounterVec

subscribesTotal *prometheus.CounterVec

messagesTotal *prometheus.CounterVec

publishedBytesTotal prometheus.Counter

receivedBytesTotal prometheus.Counter

disconnectionsTotal prometheus.Counter

connected prometheus.Gauge

latencyMeasurer *LatencyMeasurer

latencyMeasureCounter atomic.Int64

latencyErrCounter atomic.Int64

}

// BufferSize is the maximum number of unhandled messages we will buffer

// for a subscriber before dropping messages.

const BufferSize = 2048

// Subscribe calls the listener when an event matching the name is received.

func (p *PGPubsub) Subscribe(event string, listener Listener) (cancel func(), err error) {

return p.subscribeQueue(event, newMsgQueue(context.Background(), listener, nil))

}

func (p *PGPubsub) SubscribeWithErr(event string, listener ListenerWithErr) (cancel func(), err error) {

return p.subscribeQueue(event, newMsgQueue(context.Background(), nil, listener))

}

func (p *PGPubsub) subscribeQueue(event string, newQ *msgQueue) (cancel func(), err error) {

defer func() {

if err != nil {

// if we hit an error, we need to close the queue so we don't

// leak its goroutine.

newQ.close()

p.subscribesTotal.WithLabelValues("false").Inc()

} else {

p.subscribesTotal.WithLabelValues("true").Inc()

}

}()

var (

unlistenInProgress <-chan struct{}

// MUST hold the p.qMu lock to manipulate this!

qs *queueSet

)

func() {

p.qMu.Lock()

defer p.qMu.Unlock()

var ok bool

if qs, ok = p.queues[event]; !ok {

qs = newQueueSet()

p.queues[event] = qs

}

qs.m[newQ] = struct{}{}

unlistenInProgress = qs.unlistenInProgress

}()

// NOTE there cannot be any `return` statements between here and the next +-+, otherwise the

// assumptions the defer makes could be violated

if unlistenInProgress != nil {

// We have to wait here because we don't want our `Listen` call to happen before the other

// goroutine calls `Unlisten`. That would result in this subscription not getting any

// events. c.f. https://github.com/coder/coder/issues/15312

p.logger.Debug(context.Background(), "waiting for Unlisten in progress", slog.F("event", event))

<-unlistenInProgress

p.logger.Debug(context.Background(), "unlistening complete", slog.F("event", event))

}

// +-+ (see above)

defer func() {

if err != nil {

p.qMu.Lock()

defer p.qMu.Unlock()

delete(qs.m, newQ)

if len(qs.m) == 0 {

// we know that newQ was in the queueSet since we last unlocked, so there cannot

// have been any _new_ goroutines trying to Unlisten(). Therefore, if the queueSet

// is now empty, it's safe to delete.

delete(p.queues, event)

}

}

}()

// The pgListener waits for the response to `LISTEN` on a mainloop that also dispatches

// notifies. We need to avoid holding the mutex while this happens, since holding the mutex

// blocks reading notifications and can deadlock the pgListener.

// c.f. https://github.com/coder/coder/issues/11950

err = p.pgListener.Listen(event)

if err == nil {

p.logger.Debug(context.Background(), "started listening to event channel", slog.F("event", event))

}

if errors.Is(err, pq.ErrChannelAlreadyOpen) {

// It's ok if it's already open!

err = nil

}

if err != nil {

return nil, xerrors.Errorf("listen: %w", err)

}

return func() {

var unlistening chan struct{}

func() {

p.qMu.Lock()

defer p.qMu.Unlock()

newQ.close()

qSet, ok := p.queues[event]

if !ok {

p.logger.Critical(context.Background(), "event was removed before cancel", slog.F("event", event))

return

}

delete(qSet.m, newQ)

if len(qSet.m) == 0 {

unlistening = make(chan struct{})

qSet.unlistenInProgress = unlistening

}

}()

// as above, we must not hold the lock while calling into pgListener

if unlistening != nil {

uErr := p.pgListener.Unlisten(event)

close(unlistening)

// we can now delete the queueSet if it is empty.

func() {

p.qMu.Lock()

defer p.qMu.Unlock()

qSet, ok := p.queues[event]

if ok && len(qSet.m) == 0 {

p.logger.Debug(context.Background(), "removing queueSet", slog.F("event", event))

delete(p.queues, event)

}

}()

p.closeMu.Lock()

defer p.closeMu.Unlock()

if uErr != nil && !p.closedListener {

p.logger.Warn(context.Background(), "failed to unlisten", slog.Error(uErr), slog.F("event", event))

} else {

p.logger.Debug(context.Background(), "stopped listening to event channel", slog.F("event", event))

}

}

}, nil

}

func (p *PGPubsub) Publish(event string, message []byte) error {

p.logger.Debug(context.Background(), "publish", slog.F("event", event), slog.F("message_len", len(message)))

// This is safe because we are calling pq.QuoteLiteral. pg_notify doesn't

// support the first parameter being a prepared statement.

//nolint:gosec

_, err := p.db.ExecContext(context.Background(), `select pg_notify(`+pq.QuoteLiteral(event)+`, $1)`, message)

if err != nil {

p.publishesTotal.WithLabelValues("false").Inc()

return xerrors.Errorf("exec pg_notify: %w", err)

}

p.publishesTotal.WithLabelValues("true").Inc()

p.publishedBytesTotal.Add(float64(len(message)))

return nil

}

// Close closes the pubsub instance.

func (p *PGPubsub) Close() error {

p.logger.Info(context.Background(), "pubsub is closing")

err := p.closeListener()

<-p.listenDone

p.logger.Debug(context.Background(), "pubsub closed")

return err

}

// closeListener closes the pgListener, unless it has already been closed.

func (p *PGPubsub) closeListener() error {

p.closeMu.Lock()

defer p.closeMu.Unlock()

if p.closedListener {

return p.closeListenerErr

}

p.closedListener = true

p.closeListenerErr = p.pgListener.Close()

return p.closeListenerErr

}

// listen begins receiving messages on the pq listener.

func (p *PGPubsub) listen() {

defer func() {

p.logger.Info(context.Background(), "pubsub listen stopped receiving notify")

close(p.listenDone)

}()

notify := p.pgListener.NotifyChan()

for notif := range notify {

// A nil notification can be dispatched on reconnect.

if notif == nil {

p.logger.Debug(context.Background(), "notifying subscribers of a reconnection")

p.recordReconnect()

continue

}

p.listenReceive(notif)

}

}

func (p *PGPubsub) listenReceive(notif *pq.Notification) {

sizeLabel := messageSizeNormal

if len(notif.Extra) >= colossalThreshold {

sizeLabel = messageSizeColossal

}

p.messagesTotal.WithLabelValues(sizeLabel).Inc()

p.receivedBytesTotal.Add(float64(len(notif.Extra)))

p.qMu.Lock()

defer p.qMu.Unlock()

qSet, ok := p.queues[notif.Channel]

if !ok {

return

}

extra := []byte(notif.Extra)

for q := range qSet.m {

q.enqueue(extra)

}

}

func (p *PGPubsub) recordReconnect() {

p.qMu.Lock()

defer p.qMu.Unlock()

for _, qSet := range p.queues {

for q := range qSet.m {

q.dropped()

}

}

}

// logDialer is a pq.Dialer and pq.DialerContext that logs when it starts

// connecting and when the TCP connection is established.

type logDialer struct {

logger slog.Logger

d net.Dialer

}

var (

_ pq.Dialer = logDialer{}

_ pq.DialerContext = logDialer{}

)

func (d logDialer) Dial(network, address string) (net.Conn, error) {

ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)

defer cancel()

return d.DialContext(ctx, network, address)

}

func (d logDialer) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {

ctx, cancel := context.WithTimeout(context.Background(), timeout)

defer cancel()

return d.DialContext(ctx, network, address)

}

func (d logDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {

deadline, hasDeadline := ctx.Deadline()

timeoutMS := 0

if hasDeadline {

timeoutMS = int(time.Until(deadline) / time.Millisecond)

}

logger := d.logger.With(slog.F("network", network), slog.F("address", address), slog.F("timeout_ms", timeoutMS))

logger.Debug(ctx, "pubsub dialing postgres")

start := time.Now()

conn, err := d.d.DialContext(ctx, network, address)

if err != nil {

logger.Error(ctx, "pubsub failed to dial postgres")

return nil, err

}

elapsed := time.Since(start)

logger.Debug(ctx, "pubsub postgres TCP connection established", slog.F("elapsed_ms", elapsed.Milliseconds()))

return conn, nil

}

func (p *PGPubsub) startListener(ctx context.Context, connectURL string) error {

p.connected.Set(0)

// Creates a new listener using pq.

var (

dialer = logDialer{

logger: p.logger,

// pq.defaultDialer uses a zero net.Dialer as well.

d: net.Dialer{},

}

connector driver.Connector

err error

)

// Create a custom connector if the database driver supports it.

connectorCreator, ok := p.db.Driver().(database.ConnectorCreator)

if ok {

connector, err = connectorCreator.Connector(connectURL)

if err != nil {

return xerrors.Errorf("create custom connector: %w", err)

}

} else {

// use the default pq connector otherwise

connector, err = pq.NewConnector(connectURL)

if err != nil {

return xerrors.Errorf("create pq connector: %w", err)

}

}

// Set the dialer if the connector supports it.

dc, ok := connector.(database.DialerConnector)

if !ok {

p.logger.Critical(ctx, "connector does not support setting log dialer, database connection debug logs will be missing")

} else {

dc.Dialer(dialer)

}

var (

errCh = make(chan error, 1)

sentErrCh = false

)

p.pgListener = pqListenerShim{

Listener: pq.NewConnectorListener(connector, connectURL, time.Second, time.Minute, func(t pq.ListenerEventType, err error) {

switch t {

case pq.ListenerEventConnected:

p.logger.Debug(ctx, "pubsub connected to postgres")

p.connected.Set(1.0)

case pq.ListenerEventDisconnected:

p.logger.Error(ctx, "pubsub disconnected from postgres", slog.Error(err))

p.connected.Set(0)

case pq.ListenerEventReconnected:

p.logger.Info(ctx, "pubsub reconnected to postgres")

p.connected.Set(1)

case pq.ListenerEventConnectionAttemptFailed:

p.logger.Error(ctx, "pubsub failed to connect to postgres", slog.Error(err))

}

// This callback gets events whenever the connection state changes.

// Only send the first error.

if sentErrCh {

return

}

errCh <- err // won't block because we are buffered.

sentErrCh = true

}),

}

// We don't respect context cancellation here. There's a bug in the pq library

// where if you close the listener before or while the connection is being

// established, the connection will be established anyway, and will not be

// closed.

// https://github.com/lib/pq/issues/1192

if err := <-errCh; err != nil {

_ = p.pgListener.Close()

return xerrors.Errorf("create pq listener: %w", err)

}

return nil

}

// these are the metrics we compute implicitly from our existing data structures

var (

currentSubscribersDesc = prometheus.NewDesc(

"coder_pubsub_current_subscribers",

"The current number of active pubsub subscribers",

nil, nil,

)

currentEventsDesc = prometheus.NewDesc(

"coder_pubsub_current_events",

"The current number of pubsub event channels listened for",

nil, nil,

)

)

// additional metrics collected out-of-band

var (

pubsubSendLatencyDesc = prometheus.NewDesc(

"coder_pubsub_send_latency_seconds",

"The time taken to send a message into a pubsub event channel",

nil, nil,

)

pubsubRecvLatencyDesc = prometheus.NewDesc(

"coder_pubsub_receive_latency_seconds",

"The time taken to receive a message from a pubsub event channel",

nil, nil,

)

pubsubLatencyMeasureCountDesc = prometheus.NewDesc(

"coder_pubsub_latency_measures_total",

"The number of pubsub latency measurements",

nil, nil,

)

pubsubLatencyMeasureErrDesc = prometheus.NewDesc(

"coder_pubsub_latency_measure_errs_total",

"The number of pubsub latency measurement failures",

nil, nil,

)

)

// We'll track messages as size "normal" and "colossal", where the

// latter are messages larger than 7600 bytes, or 95% of the postgres

// notify limit. If we see a lot of colossal packets that's an indication that

// we might be trying to send too much data over the pubsub and are in danger of

// failing to publish.

const (

colossalThreshold = 7600

messageSizeNormal = "normal"

messageSizeColossal = "colossal"

)

// Describe implements, along with Collect, the prometheus.Collector interface

// for metrics.

func (p *PGPubsub) Describe(descs chan<- *prometheus.Desc) {

// explicit metrics

p.publishesTotal.Describe(descs)

p.subscribesTotal.Describe(descs)

p.messagesTotal.Describe(descs)

p.publishedBytesTotal.Describe(descs)

p.receivedBytesTotal.Describe(descs)

p.disconnectionsTotal.Describe(descs)

p.connected.Describe(descs)

// implicit metrics

descs <- currentSubscribersDesc

descs <- currentEventsDesc

// additional metrics

descs <- pubsubSendLatencyDesc

descs <- pubsubRecvLatencyDesc

descs <- pubsubLatencyMeasureCountDesc

descs <- pubsubLatencyMeasureErrDesc

}

// Collect implements, along with Describe, the prometheus.Collector interface

// for metrics

func (p *PGPubsub) Collect(metrics chan<- prometheus.Metric) {

// explicit metrics

p.publishesTotal.Collect(metrics)

p.subscribesTotal.Collect(metrics)

p.messagesTotal.Collect(metrics)

p.publishedBytesTotal.Collect(metrics)

p.receivedBytesTotal.Collect(metrics)

p.disconnectionsTotal.Collect(metrics)

p.connected.Collect(metrics)

// implicit metrics

p.qMu.Lock()

events := len(p.queues)

subs := 0

for _, qSet := range p.queues {

subs += len(qSet.m)

}

p.qMu.Unlock()

metrics <- prometheus.MustNewConstMetric(currentSubscribersDesc, prometheus.GaugeValue, float64(subs))

metrics <- prometheus.MustNewConstMetric(currentEventsDesc, prometheus.GaugeValue, float64(events))

// additional metrics

ctx, cancel := context.WithTimeout(context.Background(), LatencyMeasureTimeout)

defer cancel()

send, recv, err := p.latencyMeasurer.Measure(ctx, p)

metrics <- prometheus.MustNewConstMetric(pubsubLatencyMeasureCountDesc, prometheus.CounterValue, float64(p.latencyMeasureCounter.Add(1)))

if err != nil {

p.logger.Warn(context.Background(), "failed to measure latency", slog.Error(err))

metrics <- prometheus.MustNewConstMetric(pubsubLatencyMeasureErrDesc, prometheus.CounterValue, float64(p.latencyErrCounter.Add(1)))

return

}

metrics <- prometheus.MustNewConstMetric(pubsubSendLatencyDesc, prometheus.GaugeValue, send.Seconds())

metrics <- prometheus.MustNewConstMetric(pubsubRecvLatencyDesc, prometheus.GaugeValue, recv.Seconds())

}

// New creates a new Pubsub implementation using a PostgreSQL connection.

func New(startCtx context.Context, logger slog.Logger, db *sql.DB, connectURL string) (*PGPubsub, error) {

p := newWithoutListener(logger, db)

if err := p.startListener(startCtx, connectURL); err != nil {

return nil, err

}

go p.listen()

logger.Debug(startCtx, "pubsub has started")

return p, nil

}

// newWithoutListener creates a new PGPubsub without creating the pqListener.

func newWithoutListener(logger slog.Logger, db *sql.DB) *PGPubsub {

return &PGPubsub{

logger: logger,

listenDone: make(chan struct{}),

db: db,

queues: make(map[string]*queueSet),

latencyMeasurer: NewLatencyMeasurer(logger.Named("latency-measurer")),

publishesTotal: prometheus.NewCounterVec(prometheus.CounterOpts{

Namespace: "coder",

Subsystem: "pubsub",

Name: "publishes_total",

Help: "Total number of calls to Publish",

}, []string{"success"}),

subscribesTotal: prometheus.NewCounterVec(prometheus.CounterOpts{

Namespace: "coder",

Subsystem: "pubsub",

Name: "subscribes_total",

Help: "Total number of calls to Subscribe/SubscribeWithErr",

}, []string{"success"}),

messagesTotal: prometheus.NewCounterVec(prometheus.CounterOpts{

Namespace: "coder",

Subsystem: "pubsub",

Name: "messages_total",

Help: "Total number of messages received from postgres",

}, []string{"size"}),

publishedBytesTotal: prometheus.NewCounter(prometheus.CounterOpts{

Namespace: "coder",

Subsystem: "pubsub",

Name: "published_bytes_total",

Help: "Total number of bytes successfully published across all publishes",

}),

receivedBytesTotal: prometheus.NewCounter(prometheus.CounterOpts{

Namespace: "coder",

Subsystem: "pubsub",

Name: "received_bytes_total",

Help: "Total number of bytes received across all messages",

}),

disconnectionsTotal: prometheus.NewCounter(prometheus.CounterOpts{

Namespace: "coder",

Subsystem: "pubsub",

Name: "disconnections_total",

Help: "Total number of times we disconnected unexpectedly from postgres",

}),

connected: prometheus.NewGauge(prometheus.GaugeOpts{

Namespace: "coder",

Subsystem: "pubsub",

Name: "connected",

Help: "Whether we are connected (1) or not connected (0) to postgres",

}),

}

}