Add ability to load triemux from Content Store

lib/load_routes.go

@@ -9,6 +9,9 @@ import (
 	"strconv"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+
+	"github.com/jackc/pgxlisten"
 	"github.com/prometheus/client_golang/prometheus"
 
 	"github.com/alphagov/router/handlers"
@@ -110,6 +113,47 @@ func loadRoutesFromCS(pool PgxIface, mux *triemux.Mux, backends map[string]http.
 	return nil
 }
 
+func (rt *Router) listenForContentStoreUpdates(ctx context.Context) error {
+	listener := &pgxlisten.Listener{
+		Connect: func(ctx context.Context) (*pgx.Conn, error) {
+			c, err := rt.pool.Acquire(ctx)
+			if err != nil {
+				return nil, err
+			}
+			return c.Conn(), nil
+		},
+	}
+
+	listener.Handle("route_changes", pgxlisten.HandlerFunc(
+		func(ctx context.Context, notification *pgconn.Notification, conn *pgx.Conn) error {
+			rt.CsReloadChan <- true
+			return nil
+		}),
+	)
+
+	err := listener.Listen(ctx)
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (rt *Router) waitForReload() {
+	for range rt.CsReloadChan {
+		func() {
+			defer func() {
+				if r := recover(); r != nil {
+					logWarn(r)
+				}
+			}()
+
+			rt.reloadCsRoutes(rt.pool)
+		}()
+	}
+}
+
 func (rt *Router) reloadCsRoutes(pool PgxIface) {
 	var success bool
 	timer := prometheus.NewTimer(prometheus.ObserverFunc(func(v float64) {

lib/router.go

@@ -50,6 +50,7 @@ type Router struct {
 	logger            logger.Logger
 	opts              Options
 	ReloadChan        chan bool
+	CsReloadChan      chan bool
 	csMuxSampleRate   float64
 	pool              *pgxpool.Pool
 }
@@ -136,6 +137,7 @@ func NewRouter(o Options) (rt *Router, err error) {
 	}
 
 	reloadChan := make(chan bool, 1)
+	csReloadChan := make(chan bool, 1)
 	rt = &Router{
 		backends:          backends,
 		mux:               triemux.NewMux(),
@@ -144,12 +146,21 @@ func NewRouter(o Options) (rt *Router, err error) {
 		logger:            l,
 		opts:              o,
 		ReloadChan:        reloadChan,
+		CsReloadChan:      csReloadChan,
 		pool:              pool,
 		csMuxSampleRate:   csMuxSampleRate,
 	}
 
 	if csMuxSampleRate != 0.0 {
 		rt.reloadCsRoutes(pool)
+
+		go func() {
+			if err := rt.listenForContentStoreUpdates(context.Background()); err != nil {
+				logWarn(fmt.Sprintf("router: error in listenForContentStoreUpdates: %v", err))
+			}
+		}()
+
+		go rt.waitForReload()
 	}
 
 	go rt.pollAndReload()

lib/router_api.go

@@ -25,6 +25,11 @@ func NewAPIHandler(rout *Router) (api http.Handler, err error) {
 		case rout.ReloadChan <- true:
 		default:
 		}
+
+		select {
+		case rout.CsReloadChan <- true:
+		default:
+		}
 		logInfo("router: reload queued")
 		w.WriteHeader(http.StatusAccepted)
 		_, err := w.Write([]byte("Reload queued"))