diff --git a/cmd/soju/main.go b/cmd/soju/main.go index bd07336..214c0cc 100644 --- a/cmd/soju/main.go +++ b/cmd/soju/main.go @@ -293,6 +293,10 @@ func main() { log.Printf("server listening on %q", listen) } + if db, ok := db.(soju.MetricsCollectorDatabase); ok && srv.MetricsRegistry != nil { + srv.MetricsRegistry.MustRegister(db.MetricsCollector()) + } + sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP) diff --git a/db.go b/db.go index c2c845d..41dd42e 100644 --- a/db.go +++ b/db.go @@ -6,6 +6,8 @@ import ( "net/url" "strings" "time" + + "github.com/prometheus/client_golang/prometheus" ) type Database interface { @@ -28,6 +30,11 @@ type Database interface { StoreClientDeliveryReceipts(ctx context.Context, networkID int64, client string, receipts []DeliveryReceipt) error } +type MetricsCollectorDatabase interface { + Database + MetricsCollector() prometheus.Collector +} + func OpenDB(driver, source string) (Database, error) { switch driver { case "sqlite3": diff --git a/db_postgres.go b/db_postgres.go index 7db1edc..5ede40e 100644 --- a/db_postgres.go +++ b/db_postgres.go @@ -10,6 +10,8 @@ import ( "time" _ "github.com/lib/pq" + "github.com/prometheus/client_golang/prometheus" + promcollectors "github.com/prometheus/client_golang/prometheus/collectors" ) const postgresQueryTimeout = 5 * time.Second @@ -148,6 +150,10 @@ func (db *PostgresDB) Close() error { return db.db.Close() } +func (db *PostgresDB) MetricsCollector() prometheus.Collector { + return promcollectors.NewDBStatsCollector(db.db, "main") +} + func (db *PostgresDB) Stats(ctx context.Context) (*DatabaseStats, error) { ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout) defer cancel() diff --git a/db_sqlite.go b/db_sqlite.go index 1616c1e..2f4d99f 100644 --- a/db_sqlite.go +++ b/db_sqlite.go @@ -10,6 +10,8 @@ import ( "time" _ "github.com/mattn/go-sqlite3" + "github.com/prometheus/client_golang/prometheus" + promcollectors "github.com/prometheus/client_golang/prometheus/collectors" ) const sqliteQueryTimeout = 5 * time.Second @@ -238,6 +240,10 @@ func (db *SqliteDB) upgrade() error { return tx.Commit() } +func (db *SqliteDB) MetricsCollector() prometheus.Collector { + return promcollectors.NewDBStatsCollector(db.db, "main") +} + func (db *SqliteDB) Stats(ctx context.Context) (*DatabaseStats, error) { db.lock.RLock() defer db.lock.RUnlock()