├── .gitignore ├── go.mod ├── README.md ├── notifier ├── notifier_test.go ├── listener.go └── notifier.go ├── cmd └── main.go └── go.sum /.gitignore: -------------------------------------------------------------------------------- 1 | .env 2 | cli 3 | .vscode -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/brojonat/notifier 2 | 3 | go 1.22.3 4 | 5 | require github.com/jackc/pgx/v5 v5.6.0 6 | 7 | require ( 8 | github.com/jackc/pgpassfile v1.0.0 // indirect 9 | github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect 10 | github.com/jackc/puddle/v2 v2.2.1 // indirect 11 | github.com/matryer/is v1.4.1 12 | golang.org/x/crypto v0.17.0 // indirect 13 | golang.org/x/sync v0.1.0 // indirect 14 | golang.org/x/text v0.14.0 // indirect 15 | ) 16 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # notifier 2 | 3 | ## What is this? 4 | 5 | I was reading [this post](https://brandur.org/notifier) on HackerNews and decided I'd implement my own `notifier` package because I had some projects in mind that would benefit from some PubSub. So, that's what this is. I have my own blog post [here](https://brojonat.com/posts/go-postgres-listen-notify/) that goes over some of the details. This is a WIP as I build it out for a few different use cases, but feel free to drop an Issue or PR if you see something out of whack. 6 | -------------------------------------------------------------------------------- /notifier/notifier_test.go: -------------------------------------------------------------------------------- 1 | package notifier 2 | 3 | import ( 4 | "context" 5 | "log/slog" 6 | "os" 7 | "sync" 8 | "testing" 9 | 10 | "github.com/jackc/pgx/v5/pgxpool" 11 | "github.com/matryer/is" 12 | ) 13 | 14 | // NB: these tests assume you have a postgres server listening on localhost:5432 15 | // with username postgres and password postgres. You can trivially set this up 16 | // with Docker with the following: 17 | // 18 | // docker run --rm --name postgres -p 5432:5432 \ 19 | // -e POSTGRES_PASSWORD=postgres postgres 20 | 21 | func testPool(url string) (*pgxpool.Pool, error) { 22 | ctx := context.Background() 23 | pool, err := pgxpool.New(ctx, url) 24 | if err != nil { 25 | return nil, err 26 | } 27 | if err = pool.Ping(ctx); err != nil { 28 | return nil, err 29 | } 30 | return pool, nil 31 | } 32 | 33 | func TestNotifier(t *testing.T) { 34 | is := is.New(t) 35 | l := slog.New(slog.NewJSONHandler(os.Stdout, nil)) 36 | ctx := context.Background() 37 | ctx, cancel := context.WithCancel(ctx) 38 | wg := sync.WaitGroup{} 39 | pool, err := testPool("postgresql://postgres:postgres@localhost:5432") 40 | is.NoErr(err) 41 | 42 | li := NewListener(pool) 43 | err = li.Connect(ctx) 44 | is.NoErr(err) 45 | 46 | n := NewNotifier(l, li) 47 | wg.Add(1) 48 | go func() { 49 | n.Run(ctx) 50 | wg.Done() 51 | }() 52 | sub := n.Listen("foo") 53 | 54 | conn, err := pool.Acquire(ctx) 55 | wg.Add(1) 56 | go func() { 57 | <-sub.EstablishedC() 58 | conn.Exec(ctx, "select pg_notify('foo', '1')") 59 | conn.Exec(ctx, "select pg_notify('foo', '2')") 60 | conn.Exec(ctx, "select pg_notify('foo', '3')") 61 | conn.Exec(ctx, "select pg_notify('foo', '4')") 62 | conn.Exec(ctx, "select pg_notify('foo', '5')") 63 | wg.Done() 64 | }() 65 | is.NoErr(err) 66 | 67 | wg.Add(1) 68 | 69 | out := make(chan string) 70 | go func() { 71 | <-sub.EstablishedC() 72 | for i := 0; i < 5; i++ { 73 | msg := <-sub.NotificationC() 74 | out <- string(msg) 75 | } 76 | close(out) 77 | wg.Done() 78 | }() 79 | 80 | msgs := []string{} 81 | for r := range out { 82 | msgs = append(msgs, r) 83 | } 84 | is.Equal(msgs, []string{"1", "2", "3", "4", "5"}) 85 | 86 | cancel() 87 | sub.Unlisten(ctx) // uses background ctx anyway 88 | li.Close(ctx) 89 | wg.Wait() 90 | } 91 | -------------------------------------------------------------------------------- /cmd/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "fmt" 7 | "log/slog" 8 | "os" 9 | "path/filepath" 10 | "time" 11 | 12 | "github.com/brojonat/notifier/notifier" 13 | "github.com/jackc/pgx/v5/pgxpool" 14 | ) 15 | 16 | func main() { 17 | 18 | ctx := context.Background() 19 | l := getDefaultLogger(slog.LevelInfo) 20 | 21 | var url string 22 | flag.StringVar(&url, "dbhost", "", "DB host (postgresql://{user}:{password}@{hostname}/{db}?sslmode=require)") 23 | 24 | var topic string 25 | flag.StringVar(&topic, "channel", "", "a string") 26 | 27 | flag.Parse() 28 | 29 | if url == "" || topic == "" { 30 | fmt.Fprintf(os.Stderr, "missing required flag") 31 | os.Exit(1) 32 | return 33 | } 34 | 35 | // get a connection pool 36 | pool, err := pgxpool.New(ctx, url) 37 | if err != nil { 38 | fmt.Fprintf(os.Stderr, "error connection to DB: %v", err) 39 | os.Exit(1) 40 | } 41 | if err = pool.Ping(ctx); err != nil { 42 | fmt.Fprintf(os.Stderr, "error pinging DB: %v", err) 43 | os.Exit(1) 44 | } 45 | 46 | // setup the listener 47 | li := notifier.NewListener(pool) 48 | if err := li.Connect(ctx); err != nil { 49 | fmt.Fprintf(os.Stderr, "error setting up listener: %v", err) 50 | os.Exit(1) 51 | } 52 | 53 | // setup the notifier 54 | n := notifier.NewNotifier(l, li) 55 | go n.Run(ctx) 56 | 57 | // subscribe to the topic 58 | sub := n.Listen(topic) 59 | 60 | // indefinitely listen for updates 61 | go func() { 62 | <-sub.EstablishedC() 63 | for { 64 | select { 65 | case <-ctx.Done(): 66 | sub.Unlisten(ctx) 67 | fmt.Println("done listening for notifications") 68 | return 69 | case p := <-sub.NotificationC(): 70 | fmt.Printf("Got notification: %s\n", p) 71 | } 72 | } 73 | }() 74 | 75 | // unsubscribe after some time 76 | go func() { 77 | time.Sleep(20 * time.Second) 78 | sub.Unlisten(ctx) 79 | }() 80 | 81 | select {} 82 | } 83 | 84 | func getDefaultLogger(lvl slog.Level) *slog.Logger { 85 | return slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{ 86 | AddSource: true, 87 | Level: lvl, 88 | ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr { 89 | if a.Key == slog.SourceKey { 90 | source, _ := a.Value.Any().(*slog.Source) 91 | if source != nil { 92 | source.Function = "" 93 | source.File = filepath.Base(source.File) 94 | } 95 | } 96 | return a 97 | }, 98 | })) 99 | } 100 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 2 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 3 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 4 | github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= 5 | github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= 6 | github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk= 7 | github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= 8 | github.com/jackc/pgx/v5 v5.6.0 h1:SWJzexBzPL5jb0GEsrPMLIsi/3jOo7RHlzTjcAeDrPY= 9 | github.com/jackc/pgx/v5 v5.6.0/go.mod h1:DNZ/vlrUnhWCoFGxHAG8U2ljioxukquj7utPDgtQdTw= 10 | github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk= 11 | github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= 12 | github.com/matryer/is v1.4.1 h1:55ehd8zaGABKLXQUe2awZ99BD/PTc2ls+KV/dXphgEQ= 13 | github.com/matryer/is v1.4.1/go.mod h1:8I/i5uYgLzgsgEloJE1U6xx5HkBQpAZvepWuujKwMRU= 14 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 15 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 16 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 17 | github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= 18 | github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 19 | github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= 20 | github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= 21 | golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k= 22 | golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4= 23 | golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o= 24 | golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 25 | golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= 26 | golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= 27 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 28 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 29 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 30 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 31 | -------------------------------------------------------------------------------- /notifier/listener.go: -------------------------------------------------------------------------------- 1 | package notifier 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "sync" 7 | 8 | "github.com/jackc/pgx/v5/pgxpool" 9 | ) 10 | 11 | // Listener interface connects to the database and allows callers to listen to a 12 | // particular topic by issuing a LISTEN command. WaitForNotification blocks 13 | // until receiving a notification or until the supplied context expires. The 14 | // default implementation is tightly coupled to pgx (following River's 15 | // implementation), but callers may implement their own listeners for any 16 | // backend they'd like. 17 | type Listener interface { 18 | Close(ctx context.Context) error 19 | Connect(ctx context.Context) error 20 | Listen(ctx context.Context, topic string) error 21 | Ping(ctx context.Context) error 22 | Unlisten(ctx context.Context, topic string) error 23 | WaitForNotification(ctx context.Context) (*Notification, error) 24 | } 25 | 26 | // NewListener return a Listener that draws a connection from the supplied Pool. This 27 | // is somewhat discouraged 28 | func NewListener(dbp *pgxpool.Pool) Listener { 29 | return &listener{ 30 | mu: sync.Mutex{}, 31 | dbPool: dbp, 32 | } 33 | } 34 | 35 | type listener struct { 36 | conn *pgxpool.Conn 37 | dbPool *pgxpool.Pool 38 | mu sync.Mutex 39 | } 40 | 41 | // Close the connection to the database. 42 | func (l *listener) Close(ctx context.Context) error { 43 | l.mu.Lock() 44 | defer l.mu.Unlock() 45 | 46 | if l.conn == nil { 47 | return nil 48 | } 49 | 50 | // Release below would take care of cleanup and potentially put the 51 | // connection back into rotation, but in case a Listen was invoked without a 52 | // subsequent Unlisten on the same topic, close the connection explicitly to 53 | // guarantee no other caller will receive a partially tainted connection. 54 | err := l.conn.Conn().Close(ctx) 55 | 56 | // Even in the event of an error, make sure conn is set back to nil so that 57 | // the listener can be reused. 58 | l.conn.Release() 59 | l.conn = nil 60 | 61 | return err 62 | } 63 | 64 | // Connect to the database. 65 | func (l *listener) Connect(ctx context.Context) error { 66 | l.mu.Lock() 67 | defer l.mu.Unlock() 68 | 69 | if l.conn != nil { 70 | return errors.New("connection already established") 71 | } 72 | 73 | conn, err := l.dbPool.Acquire(ctx) 74 | if err != nil { 75 | return err 76 | } 77 | 78 | l.conn = conn 79 | return nil 80 | } 81 | 82 | // Listen issues a LISTEN command for the supplied topic. 83 | func (l *listener) Listen(ctx context.Context, topic string) error { 84 | l.mu.Lock() 85 | defer l.mu.Unlock() 86 | 87 | _, err := l.conn.Exec(ctx, "LISTEN \""+topic+"\"") 88 | return err 89 | } 90 | 91 | // Ping the database 92 | func (l *listener) Ping(ctx context.Context) error { 93 | l.mu.Lock() 94 | defer l.mu.Unlock() 95 | 96 | return l.conn.Ping(ctx) 97 | } 98 | 99 | // Unlisten issues an UNLISTEN from the supplied topic. 100 | func (l *listener) Unlisten(ctx context.Context, topic string) error { 101 | l.mu.Lock() 102 | defer l.mu.Unlock() 103 | 104 | _, err := l.conn.Exec(ctx, "UNLISTEN \""+topic+"\"") 105 | return err 106 | } 107 | 108 | // WaitForNotification blocks until receiving a notification and returns it. The 109 | // pgx driver should maintain a buffer of notifications, so as long as Listen 110 | // has been called, repeatedly calling WaitForNotification should yield all 111 | // notifications. 112 | func (l *listener) WaitForNotification(ctx context.Context) (*Notification, error) { 113 | l.mu.Lock() 114 | defer l.mu.Unlock() 115 | 116 | pgn, err := l.conn.Conn().WaitForNotification(ctx) 117 | 118 | if err != nil { 119 | return nil, err 120 | } 121 | 122 | n := Notification{ 123 | Channel: pgn.Channel, 124 | Payload: []byte(pgn.Payload), 125 | } 126 | 127 | return &n, nil 128 | } 129 | -------------------------------------------------------------------------------- /notifier/notifier.go: -------------------------------------------------------------------------------- 1 | package notifier 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "log/slog" 8 | "slices" 9 | "sync" 10 | "time" 11 | ) 12 | 13 | // Notifier interface wraps a Listener. It holds a single Postgres connection 14 | // per process, allows other components in the same program to use it to 15 | // subscribe to any number of topics, waits for notifications, and distributes 16 | // them to listening components as they’re received 17 | type Notifier interface { 18 | // Returns a Subscription to the supplied channel topic which can be used to by 19 | // the caller to receive data published to that channel 20 | Listen(channel string) Subscription 21 | 22 | // this runs the receiving loop forever 23 | Run(ctx context.Context) error 24 | } 25 | 26 | // Subscription provides a means to listen on a particular topic. Notifiers 27 | // return Subscriptions that callers can use to receive updates. 28 | type Subscription interface { 29 | NotificationC() <-chan []byte 30 | EstablishedC() <-chan struct{} 31 | Unlisten(ctx context.Context) 32 | } 33 | 34 | // Notification encapsulates a published message 35 | type Notification struct { 36 | Channel string `json:"channel"` 37 | Payload []byte `json:"payload"` 38 | } 39 | 40 | type subscription struct { 41 | channel string 42 | listenChan chan []byte 43 | notifier *notifier 44 | 45 | establishedChan chan struct{} 46 | establishedChanClose func() 47 | unlistenOnce sync.Once 48 | } 49 | 50 | // NotificationC returns the underlying notification channel. 51 | func (s *subscription) NotificationC() <-chan []byte { return s.listenChan } 52 | 53 | // EstablishedC is a channel that's closed after the Notifier has successfully 54 | // established a connection to the database and started listening for updates. 55 | // 56 | // There's no full guarantee that the notifier will successfully establish a 57 | // listen, so callers will usually want to `select` on it combined with a 58 | // context done, a stop channel, and/or a timeout. 59 | func (s *subscription) EstablishedC() <-chan struct{} { return s.establishedChan } 60 | 61 | // Unlisten unregisters the subscriber from its notifier 62 | func (s *subscription) Unlisten(ctx context.Context) { 63 | s.unlistenOnce.Do(func() { 64 | // Unlisten uses background context in case of cancellation. 65 | if err := s.notifier.unlisten(context.Background(), s); err != nil { 66 | s.notifier.logger.Error("error unlistening on channel", "err", err, "channel", s.channel) 67 | } 68 | }) 69 | } 70 | 71 | type notifier struct { 72 | mu sync.RWMutex 73 | logger *slog.Logger 74 | listener Listener 75 | subscriptions map[string][]*subscription 76 | channelChanges []channelChange 77 | waitForNotificationCancel context.CancelFunc 78 | } 79 | 80 | func NewNotifier(l *slog.Logger, li Listener) Notifier { 81 | return ¬ifier{ 82 | mu: sync.RWMutex{}, 83 | logger: l, 84 | listener: li, 85 | subscriptions: make(map[string][]*subscription), 86 | channelChanges: []channelChange{}, 87 | waitForNotificationCancel: context.CancelFunc(func() {}), 88 | } 89 | } 90 | 91 | type channelChange struct { 92 | channel string 93 | close func() 94 | operation string 95 | } 96 | 97 | // Listen returns a Subscription. 98 | func (n *notifier) Listen(channel string) Subscription { 99 | n.mu.Lock() 100 | defer n.mu.Unlock() 101 | 102 | existingSubs := n.subscriptions[channel] 103 | 104 | sub := &subscription{ 105 | channel: channel, 106 | listenChan: make(chan []byte, 2), 107 | notifier: n, 108 | } 109 | n.subscriptions[channel] = append(existingSubs, sub) 110 | 111 | if len(existingSubs) > 0 { 112 | // If there's already another subscription for this channel, reuse its 113 | // established channel. It may already be closed (to indicate that the 114 | // connection is established), but that's okay. 115 | sub.establishedChan = existingSubs[0].establishedChan 116 | sub.establishedChanClose = func() {} // no op since not channel owner 117 | 118 | return sub 119 | } 120 | 121 | // The notifier will close this channel after it's successfully established 122 | // `LISTEN` for the given channel. Gives subscribers a way to confirm a 123 | // listen before moving on, which is especially useful in tests. 124 | sub.establishedChan = make(chan struct{}) 125 | sub.establishedChanClose = sync.OnceFunc(func() { close(sub.establishedChan) }) 126 | 127 | n.channelChanges = append(n.channelChanges, 128 | channelChange{channel, sub.establishedChanClose, "listen"}) 129 | 130 | // Cancel out of blocking on WaitForNotification so changes can be processed 131 | // immediately. 132 | n.waitForNotificationCancel() 133 | 134 | return sub 135 | } 136 | 137 | const listenerTimeout = 10 * time.Second 138 | 139 | // Listens on a topic with an appropriate logging statement. Should be preferred 140 | // to `listener.Listen` for improved logging/telemetry. 141 | func (n *notifier) listenerListen(ctx context.Context, channel string) error { 142 | ctx, cancel := context.WithTimeout(ctx, listenerTimeout) 143 | defer cancel() 144 | 145 | n.logger.Debug("listening on channel", "channel", channel) 146 | if err := n.listener.Listen(ctx, channel); err != nil { 147 | return fmt.Errorf("error listening on channel %q: %w", channel, err) 148 | } 149 | 150 | return nil 151 | } 152 | 153 | // Unlistens on a topic with an appropriate logging statement. Should be 154 | // preferred to `listener.Unlisten` for improved logging/telemetry. 155 | func (n *notifier) listenerUnlisten(ctx context.Context, channel string) error { 156 | ctx, cancel := context.WithTimeout(ctx, listenerTimeout) 157 | defer cancel() 158 | 159 | n.logger.Debug("unlistening on channel", "channel", channel) 160 | if err := n.listener.Unlisten(ctx, string(channel)); err != nil { 161 | return fmt.Errorf("error unlistening on channel %q: %w", channel, err) 162 | } 163 | 164 | return nil 165 | } 166 | 167 | // this needs to pull channelChange instances from the channelChange channel 168 | // in order to perform LISTEN/UNLISTEN operations on the notifier. 169 | func (n *notifier) processChannelChanges(ctx context.Context) error { 170 | n.logger.Debug("processing channel changes...") 171 | n.mu.Lock() 172 | defer n.mu.Unlock() 173 | for _, u := range n.channelChanges { 174 | switch u.operation { 175 | case "listen": 176 | n.logger.Debug("listening to new channel", "channel", u.channel) 177 | n.listenerListen(ctx, u.channel) 178 | u.close() 179 | case "unlisten": 180 | n.logger.Debug("unlistening from channel", "channel", u.channel) 181 | n.listenerUnlisten(ctx, u.channel) 182 | default: 183 | n.logger.Error("got unexpected change operation", "operation", u.operation) 184 | } 185 | } 186 | return nil 187 | } 188 | 189 | // waitOnce blocks until either 1) a notification is received and 190 | // distributed to all topic listeners, 2) the timeout is hit, or 3) an external 191 | // caller calls l.waitForNotificationCancel. In all 3 cases, nil is returned to 192 | // signal good/expected exit conditions, meaning a caller can simply call 193 | // handleNextNotification again. 194 | func (n *notifier) waitOnce(ctx context.Context) error { 195 | if err := n.processChannelChanges(ctx); err != nil { 196 | return err 197 | } 198 | 199 | // WaitForNotification is a blocking function, but since we want to wake 200 | // occasionally to process new `LISTEN`/`UNLISTEN`, we let the context 201 | // timeout and also expose a way for external code to cancel this loop with 202 | // waitForNotificationCancel. 203 | notification, err := func() (*Notification, error) { 204 | const listenTimeout = 30 * time.Second 205 | 206 | ctx, cancel := context.WithTimeout(ctx, listenTimeout) 207 | defer cancel() 208 | 209 | // Provides a way for the blocking wait to be cancelled in case a new 210 | // subscription change comes in. 211 | n.mu.Lock() 212 | n.waitForNotificationCancel = cancel 213 | n.mu.Unlock() 214 | 215 | notification, err := n.listener.WaitForNotification(ctx) 216 | if err != nil { 217 | return nil, fmt.Errorf("error waiting for notification: %w", err) 218 | } 219 | 220 | return notification, nil 221 | }() 222 | if err != nil { 223 | // If the error was a cancellation or the deadline being exceeded but 224 | // there's no error in the parent context, return no error. 225 | if (errors.Is(err, context.Canceled) || 226 | errors.Is(err, context.DeadlineExceeded)) && ctx.Err() == nil { 227 | return nil 228 | } 229 | 230 | return err 231 | } 232 | 233 | n.mu.RLock() 234 | defer n.mu.RUnlock() 235 | 236 | for _, sub := range n.subscriptions[notification.Channel] { 237 | select { 238 | case sub.listenChan <- []byte(notification.Payload): 239 | default: 240 | n.logger.Error("dropped notification due to full buffer", "payload", notification.Payload) 241 | } 242 | } 243 | 244 | return nil 245 | } 246 | 247 | func (n *notifier) unlisten(ctx context.Context, sub *subscription) error { 248 | n.mu.Lock() 249 | defer n.mu.Unlock() 250 | 251 | subs := n.subscriptions[sub.channel] 252 | 253 | // stop listening if last subscriber 254 | if len(subs) <= 1 { 255 | // UNLISTEN for this channel 256 | n.listenerUnlisten(ctx, sub.channel) 257 | } 258 | 259 | // remove subscription from the subscriptions map 260 | n.subscriptions[sub.channel] = slices.DeleteFunc(n.subscriptions[sub.channel], func(s *subscription) bool { 261 | return s == sub 262 | }) 263 | if len(n.subscriptions[sub.channel]) < 1 { 264 | delete(n.subscriptions, sub.channel) 265 | } 266 | n.logger.Debug("removed subscription", "new_num_subscriptions", len(n.subscriptions[sub.channel]), "channel", sub.channel) 267 | 268 | return nil 269 | } 270 | 271 | func (n *notifier) Run(ctx context.Context) error { 272 | for { 273 | err := n.waitOnce(ctx) 274 | if err != nil || ctx.Err() != nil { 275 | return err 276 | } 277 | } 278 | } 279 | --------------------------------------------------------------------------------