path: root/src/go/collectors/go.d.plugin/modules/pgbouncer/collect.go

// SPDX-License-Identifier: GPL-3.0-or-later

package pgbouncer

import (
	"context"
	"database/sql"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/blang/semver/v4"
	"github.com/jackc/pgx/v4"
	"github.com/jackc/pgx/v4/stdlib"
)

// 'SHOW STATS;' response was changed significantly in v1.8.0
// v1.8.0 was released in 2015 - no need to complicate the code to support the old version.
var minSupportedVersion = semver.Version{Major: 1, Minor: 8, Patch: 0}

const (
	queryShowVersion   = "SHOW VERSION;"
	queryShowConfig    = "SHOW CONFIG;"
	queryShowDatabases = "SHOW DATABASES;"
	queryShowStats     = "SHOW STATS;"
	queryShowPools     = "SHOW POOLS;"
)

func (p *PgBouncer) collect() (map[string]int64, error) {
	if p.db == nil {
		if err := p.openConnection(); err != nil {
			return nil, err
		}
	}
	if p.version == nil {
		ver, err := p.queryVersion()
		if err != nil {
			return nil, err
		}
		p.Debugf("connected to PgBouncer v%s", ver)
		if ver.LE(minSupportedVersion) {
			return nil, fmt.Errorf("unsupported version: v%s, required v%s+", ver, minSupportedVersion)
		}
		p.version = ver
	}

	now := time.Now()
	if now.Sub(p.recheckSettingsTime) > p.recheckSettingsEvery {
		v, err := p.queryMaxClientConn()
		if err != nil {
			return nil, err
		}
		p.maxClientConn = v
	}

	// http://www.pgbouncer.org/usage.html

	p.resetMetrics()

	if err := p.collectDatabases(); err != nil {
		return nil, err
	}
	if err := p.collectStats(); err != nil {
		return nil, err
	}
	if err := p.collectPools(); err != nil {
		return nil, err
	}

	mx := make(map[string]int64)
	p.collectMetrics(mx)

	return mx, nil
}

func (p *PgBouncer) collectMetrics(mx map[string]int64) {
	var clientConns int64
	for name, db := range p.metrics.dbs {
		if !db.updated {
			delete(p.metrics.dbs, name)
			p.removeDatabaseCharts(name)
			continue
		}
		if !db.hasCharts {
			db.hasCharts = true
			p.addNewDatabaseCharts(name, db.pgDBName)
		}

		mx["db_"+name+"_total_xact_count"] = db.totalXactCount
		mx["db_"+name+"_total_xact_time"] = db.totalXactTime
		mx["db_"+name+"_avg_xact_time"] = db.avgXactTime

		mx["db_"+name+"_total_query_count"] = db.totalQueryCount
		mx["db_"+name+"_total_query_time"] = db.totalQueryTime
		mx["db_"+name+"_avg_query_time"] = db.avgQueryTime

		mx["db_"+name+"_total_wait_time"] = db.totalWaitTime
		mx["db_"+name+"_maxwait"] = db.maxWait*1e6 + db.maxWaitUS

		mx["db_"+name+"_cl_active"] = db.clActive
		mx["db_"+name+"_cl_waiting"] = db.clWaiting
		mx["db_"+name+"_cl_cancel_req"] = db.clCancelReq
		clientConns += db.clActive + db.clWaiting + db.clCancelReq

		mx["db_"+name+"_sv_active"] = db.svActive
		mx["db_"+name+"_sv_idle"] = db.svIdle
		mx["db_"+name+"_sv_used"] = db.svUsed
		mx["db_"+name+"_sv_tested"] = db.svTested
		mx["db_"+name+"_sv_login"] = db.svLogin

		mx["db_"+name+"_total_received"] = db.totalReceived
		mx["db_"+name+"_total_sent"] = db.totalSent

		mx["db_"+name+"_sv_conns_utilization"] = calcPercentage(db.currentConnections, db.maxConnections)
	}

	mx["cl_conns_utilization"] = calcPercentage(clientConns, p.maxClientConn)
}

func (p *PgBouncer) collectDatabases() error {
	q := queryShowDatabases
	p.Debugf("executing query: %v", q)

	var db string
	return p.collectQuery(q, func(column, value string) {
		switch column {
		case "name":
			db = value
			p.getDBMetrics(db).updated = true
		case "database":
			p.getDBMetrics(db).pgDBName = value
		case "max_connections":
			p.getDBMetrics(db).maxConnections = parseInt(value)
		case "current_connections":
			p.getDBMetrics(db).currentConnections = parseInt(value)
		case "paused":
			p.getDBMetrics(db).paused = parseInt(value)
		case "disabled":
			p.getDBMetrics(db).disabled = parseInt(value)
		}
	})
}

func (p *PgBouncer) collectStats() error {
	q := queryShowStats
	p.Debugf("executing query: %v", q)

	var db string
	return p.collectQuery(q, func(column, value string) {
		switch column {
		case "database":
			db = value
			p.getDBMetrics(db).updated = true
		case "total_xact_count":
			p.getDBMetrics(db).totalXactCount = parseInt(value)
		case "total_query_count":
			p.getDBMetrics(db).totalQueryCount = parseInt(value)
		case "total_received":
			p.getDBMetrics(db).totalReceived = parseInt(value)
		case "total_sent":
			p.getDBMetrics(db).totalSent = parseInt(value)
		case "total_xact_time":
			p.getDBMetrics(db).totalXactTime = parseInt(value)
		case "total_query_time":
			p.getDBMetrics(db).totalQueryTime = parseInt(value)
		case "total_wait_time":
			p.getDBMetrics(db).totalWaitTime = parseInt(value)
		case "avg_xact_time":
			p.getDBMetrics(db).avgXactTime = parseInt(value)
		case "avg_query_time":
			p.getDBMetrics(db).avgQueryTime = parseInt(value)
		}
	})
}

func (p *PgBouncer) collectPools() error {
	q := queryShowPools
	p.Debugf("executing query: %v", q)

	// an entry is made for each couple of (database, user).
	var db string
	return p.collectQuery(q, func(column, value string) {
		switch column {
		case "database":
			db = value
			p.getDBMetrics(db).updated = true
		case "cl_active":
			p.getDBMetrics(db).clActive += parseInt(value)
		case "cl_waiting":
			p.getDBMetrics(db).clWaiting += parseInt(value)
		case "cl_cancel_req":
			p.getDBMetrics(db).clCancelReq += parseInt(value)
		case "sv_active":
			p.getDBMetrics(db).svActive += parseInt(value)
		case "sv_idle":
			p.getDBMetrics(db).svIdle += parseInt(value)
		case "sv_used":
			p.getDBMetrics(db).svUsed += parseInt(value)
		case "sv_tested":
			p.getDBMetrics(db).svTested += parseInt(value)
		case "sv_login":
			p.getDBMetrics(db).svLogin += parseInt(value)
		case "maxwait":
			p.getDBMetrics(db).maxWait += parseInt(value)
		case "maxwait_us":
			p.getDBMetrics(db).maxWaitUS += parseInt(value)
		}
	})
}

func (p *PgBouncer) queryMaxClientConn() (int64, error) {
	q := queryShowConfig
	p.Debugf("executing query: %v", q)

	var v int64
	var key string
	err := p.collectQuery(q, func(column, value string) {
		switch column {
		case "key":
			key = value
		case "value":
			if key == "max_client_conn" {
				v = parseInt(value)
			}
		}
	})
	return v, err
}

var reVersion = regexp.MustCompile(`\d+\.\d+\.\d+`)

func (p *PgBouncer) queryVersion() (*semver.Version, error) {
	q := queryShowVersion
	p.Debugf("executing query: %v", q)

	var resp string
	ctx, cancel := context.WithTimeout(context.Background(), p.Timeout.Duration())
	defer cancel()
	if err := p.db.QueryRowContext(ctx, q).Scan(&resp); err != nil {
		return nil, err
	}

	if !strings.Contains(resp, "PgBouncer") {
		return nil, fmt.Errorf("not PgBouncer instance: version response: %s", resp)
	}

	ver := reVersion.FindString(resp)
	if ver == "" {
		return nil, fmt.Errorf("couldn't parse version string '%s' (expected pattern '%s')", resp, reVersion)
	}

	v, err := semver.New(ver)
	if err != nil {
		return nil, fmt.Errorf("couldn't parse version string '%s': %v", ver, err)
	}

	return v, nil
}

func (p *PgBouncer) openConnection() error {
	cfg, err := pgx.ParseConfig(p.DSN)
	if err != nil {
		return err
	}
	cfg.PreferSimpleProtocol = true

	db, err := sql.Open("pgx", stdlib.RegisterConnConfig(cfg))
	if err != nil {
		return fmt.Errorf("error on opening a connection with the PgBouncer database [%s]: %v", p.DSN, err)
	}

	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	db.SetConnMaxLifetime(10 * time.Minute)

	p.db = db

	return nil
}

func (p *PgBouncer) collectQuery(query string, assign func(column, value string)) error {
	ctx, cancel := context.WithTimeout(context.Background(), p.Timeout.Duration())
	defer cancel()
	rows, err := p.db.QueryContext(ctx, query)
	if err != nil {
		return err
	}
	defer func() { _ = rows.Close() }()

	columns, err := rows.Columns()
	if err != nil {
		return err
	}

	values := makeNullStrings(len(columns))
	for rows.Next() {
		if err := rows.Scan(values...); err != nil {
			return err
		}
		for i, v := range values {
			assign(columns[i], valueToString(v))
		}
	}
	return rows.Err()
}

func (p *PgBouncer) getDBMetrics(dbname string) *dbMetrics {
	db, ok := p.metrics.dbs[dbname]
	if !ok {
		db = &dbMetrics{name: dbname}
		p.metrics.dbs[dbname] = db
	}
	return db
}

func (p *PgBouncer) resetMetrics() {
	for name, db := range p.metrics.dbs {
		p.metrics.dbs[name] = &dbMetrics{
			name:      db.name,
			pgDBName:  db.pgDBName,
			hasCharts: db.hasCharts,
		}
	}
}

func valueToString(value any) string {
	v, ok := value.(*sql.NullString)
	if !ok || !v.Valid {
		return ""
	}
	return v.String
}

func makeNullStrings(size int) []any {
	vs := make([]any, size)
	for i := range vs {
		vs[i] = &sql.NullString{}
	}
	return vs
}

func parseInt(s string) int64 {
	v, _ := strconv.ParseInt(s, 10, 64)
	return v
}

func calcPercentage(value, total int64) int64 {
	if total == 0 {
		return 0
	}
	return value * 100 / total
}