github.com/decred/dcrlnd@v0.7.6/kvdb/postgres/db_conn_set.go (about)

     1  //go:build kvdb_postgres
     2  // +build kvdb_postgres
     3  
     4  package postgres
     5  
     6  import (
     7  	"database/sql"
     8  	"fmt"
     9  	"sync"
    10  
    11  	_ "github.com/jackc/pgx/v4/stdlib"
    12  )
    13  
    14  // dbConn stores the actual connection and a user count.
    15  type dbConn struct {
    16  	db    *sql.DB
    17  	count int
    18  }
    19  
    20  // dbConnSet stores a set of connections.
    21  type dbConnSet struct {
    22  	dbConn         map[string]*dbConn
    23  	maxConnections int
    24  
    25  	sync.Mutex
    26  }
    27  
    28  // newDbConnSet initializes a new set of connections.
    29  func newDbConnSet(maxConnections int) *dbConnSet {
    30  	return &dbConnSet{
    31  		dbConn:         make(map[string]*dbConn),
    32  		maxConnections: maxConnections,
    33  	}
    34  }
    35  
    36  // Open opens a new database connection. If a connection already exists for the
    37  // given dsn, the existing connection is returned.
    38  func (d *dbConnSet) Open(dsn string) (*sql.DB, error) {
    39  	d.Lock()
    40  	defer d.Unlock()
    41  
    42  	if dbConn, ok := d.dbConn[dsn]; ok {
    43  		dbConn.count++
    44  
    45  		return dbConn.db, nil
    46  	}
    47  
    48  	db, err := sql.Open("pgx", dsn)
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  
    53  	// Limit maximum number of open connections. This is useful to prevent
    54  	// the server from running out of connections and returning an error.
    55  	// With this client-side limit in place, lnd will wait for a connection
    56  	// to become available.
    57  	if d.maxConnections != 0 {
    58  		db.SetMaxOpenConns(d.maxConnections)
    59  	}
    60  
    61  	d.dbConn[dsn] = &dbConn{
    62  		db:    db,
    63  		count: 1,
    64  	}
    65  
    66  	return db, nil
    67  }
    68  
    69  // Close closes the connection with the given dsn. If there are still other
    70  // users of the same connection, this function does nothing.
    71  func (d *dbConnSet) Close(dsn string) error {
    72  	d.Lock()
    73  	defer d.Unlock()
    74  
    75  	dbConn, ok := d.dbConn[dsn]
    76  	if !ok {
    77  		return fmt.Errorf("connection not found: %v", dsn)
    78  	}
    79  
    80  	// Reduce user count.
    81  	dbConn.count--
    82  
    83  	// Do not close if there are other users.
    84  	if dbConn.count > 0 {
    85  		return nil
    86  	}
    87  
    88  	// Close connection.
    89  	delete(d.dbConn, dsn)
    90  
    91  	return dbConn.db.Close()
    92  }