"conn busy" with wrapper that handles multiple LISTEN/NOTIFY

[kohenkatz]

In my application, I need several goroutines to LISTEN for several NOTIFY channels.

Based on the discussion in #735, #195, and a few other places, here's what I wrote:

type DBClient struct {
	pool *pgxpool.Pool

	listenLock     sync.Mutex
	listenConn     *pgxpool.Conn
	listenCancel   context.CancelFunc
	listenChannels map[string][]chan string
}

func New(connStr string, logger *zap.Logger) (*DBClient, error) {
	config, err := pgxpool.ParseConfig(connStr)
	if err != nil {
		return nil, err
	}
	config.ConnConfig.Tracer = &tracelog.TraceLog{
		Logger:   pgxzap.NewLogger(logger.Named("postgres")),
		LogLevel: tracelog.LogLevelDebug,
	}

	pool, err := pgxpool.NewWithConfig(context.Background(), config)
	if err != nil {
		return nil, err
	}

	// Make sure the connection opens
	err = pool.Ping(context.Background())
	if err != nil {
		return nil, err
	}

	client := &DBClient{
		pool:           pool,
		listenChannels: make(map[string][]chan string),
	}

	return client, err
}

func (db *DBClient) startListener(ctx context.Context) error {
	// If we have no listeners, we don't need to listen to anything,
	// so clean up any resources we used
	if len(db.listenChannels) == 0 {
		// IMPORTANT: This uses the same lock, so it must be
		// called before our call to `Lock`.
		db.StopListening()
		return nil
	}

	// NOTE: This lock *does not* remain held within the goroutine,
	// only within the setup portion of the function.
	db.listenLock.Lock()
	defer db.listenLock.Unlock()

	if db.listenConn == nil || db.listenConn.Conn().IsClosed() {
		conn, err := db.pool.Acquire(ctx)
		if err != nil {
			return err
		}
		db.listenConn = conn
	}

	cCtx, cCancel := context.WithCancel(ctx)
	db.listenCancel = cCancel

	rawConn := db.listenConn.Conn()

	// Listen for all subscribed channel named
	for name := range db.listenChannels {
		if _, err := rawConn.Exec(ctx, "LISTEN "+name); err != nil {
			return err
		}
	}

	go func() {
		for {
			msg, err := rawConn.WaitForNotification(cCtx)
			if err != nil {
				if errors.Is(err, context.Canceled) {
					// We've been told to stop listening,
					// so exit the goroutine
					return
				}
				// TODO: what other types of errors might we get?
				// TODO: which types of errors can we reconnect and recover from?
				// TODO: do we need to log here, or has the trace logger already done that for us?
				continue
			}
			if msg != nil {
				if channels, ok := db.listenChannels[msg.Channel]; ok {
					for _, ch := range channels {
						ch <- msg.Payload
					}
				}
			}
		}
	}()

	return nil
}

func (db *DBClient) Subscribe(ctx context.Context, name string) (chan string, error) {
	ch := make(chan string)
	db.listenChannels[name] = append(db.listenChannels[name], ch)

	db.startListener(ctx)

	return ch, nil
}


func (db *DBClient) StopListening() {
	db.listenLock.Lock()
	defer db.listenLock.Unlock()
	// ORDER MATTERS
	// 1. Stop processing
	if db.listenCancel != nil {
		db.listenCancel()
	}
	// 2. Stop listening
	// 3. Release connection back to pool
	if db.listenConn != nil {
		db.listenConn.Exec(context.Background(), "UNLISTEN *")
		db.listenConn.Release()
	}
	// 4. Close all the listening channels
	for name, channels := range db.listenChannels {
		for _, ch := range channels {
			close(ch)
		}
		delete(db.listenChannels, name)
	}
	// 5. Clean up
	db.listenCancel = nil
	db.listenConn = nil
}

func (db *DBClient) Close() {
	db.StopListening()
	db.pool.Close()
}

// implement interface pgxscan.Querier
func (db *DBClient) Query(ctx context.Context, query string, args ...interface{}) (pgx.Rows, error) {
	return db.pool.Query(ctx, query, args...)
}

Then I call it like this:

lMedia, err := db.Subscribe(ctx, "notify_live_media")
if err != nil {
	return nil, err
}
lMsgs, err := db.Subscribe(ctx, "notify_live_messages")
if err != nil {
	return nil, err
}

The problem is that the second call to Subscribe throws a conn busy error. From the logs, I can see that the first call to Subscribe properly calls LISTEN notify_live_media, and the second call the Subscribe properly cancels the listener, then calls LISTEN notify_live_media which works properly followed by LISTEN notify_live_messages which returns the error.

As far as I can tell, the calls to rawConn.Exec() are being executed sequentially, and I see in the source that Exec calls execSimpleProtocol() which both locks the connection (inside pgConn.Exec()) and also unlocks the connection (inside mrr.Close()).

Why can't I run multiple LISTEN queries in a loop? What am I missing?

[jackc]

I didn't notice anything for sure, but I suspect there is some sort of race condition.

1. Run with the race detector.
2. Use dlv to set a breakpoint on changes to the internal lock variable that is triggering the conn busy. This might help find what is holding the lock.

---

Beyond that though, I would recommend separating the listen to channel from wait for notification steps. I think it is that coupling that makes this hard to reason about. Even simpler would be to avoid adding listen channels after the process is started. That is the approach I took for my own listener (https://github.com/jackc/pgxlisten -- I use in it production -- but not publicized because I want freedom to change or abandon it at will without worrying about the impact to others -- so maybe take more as source of ideas than importing it directly).

[bluekeyes]

I recently had a similar conn busy error with LISTEN/NOTIFY, and in my case the problem was an application-level race condition that happened when cancelling the context passed to WaitForNotification.

When you cancel that context, a pgx-internal goroutine needs to detect it, set a deadline on the underlying connection, and then Postgres has to respond to the deadline. Until WaitForNotification returns, the connection could still be active.

My situation was something like this, where I was trying to combine NOTIFY events with an external event source:

func listenUntilError(ctx context.Context, conn *pgx.Conn, external chan struct{}) error {
	for {
		event, err := waitForEvent(ctx, conn, external)
		if err != nil {
			return err
		}
		// send event to subscribers
	}
}

func waitForEvent(ctx context.Context, conn *pgx.Conn, external chan struct{}) (Event, error) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel() // cancel WaitForNotification if we get an external event

	type waitResult struct {
		n   *pgconn.Notification
		err error
	}

	result := make(chan waitResult, 1)
	go func() {
		n, err := conn.WaitForNotification(ctx)
		result <- waitResult{n, err}
	}()

	select {
	case <-external:
		return Event{Payload: "external"}, nil

	case r := <-result:
		if r.err != nil {
			return Event{}, r.err
		}
		return Event{Payload: r.n.Payload}, nil
	}
}

If an external event arrived first, I cancelled the context, but did not wait for WaitForNotification to return. Depending on how long that took, the next round of the loop in listenUntilError could try to wait on the same connection while there was still an active query.

I fixed the issue by combining the two functions in a way that removed the need to cancel the context. There's now one loop in a goroutine calling WaitForNotification and a second loop around the select that combines the events from different sources.

I'm not sure if this is the same problem, but the code in the original post also cancels the context in a WaitForNotification call without waiting for the return value.