github.com/pf-qiu/concourse/v6@v6.7.3-0.20201207032516-1f455d73275f/atc/db/connection_retrying_driver.go (about)

     1  package db
     2  
     3  import (
     4  	"database/sql"
     5  	"database/sql/driver"
     6  
     7  	"github.com/cenkalti/backoff"
     8  	"github.com/lib/pq"
     9  )
    10  
    11  type connectionRetryingDriver struct {
    12  	driver.Driver
    13  	driverName string
    14  }
    15  
    16  func SetupConnectionRetryingDriver(
    17  	delegateDriverName string,
    18  	sqlDataSource string,
    19  	newDriverName string,
    20  ) {
    21  	for _, driverName := range sql.Drivers() {
    22  		if driverName == newDriverName {
    23  			return
    24  		}
    25  	}
    26  	delegateDBConn, err := sql.Open(delegateDriverName, sqlDataSource)
    27  	if err == nil {
    28  		// ignoring any connection errors since we only need this to access the driver struct
    29  		_ = delegateDBConn.Close()
    30  	}
    31  
    32  	connectionRetryingDriver := &connectionRetryingDriver{
    33  		delegateDBConn.Driver(),
    34  		delegateDriverName,
    35  	}
    36  	sql.Register(newDriverName, connectionRetryingDriver)
    37  }
    38  
    39  func (d *connectionRetryingDriver) Open(name string) (driver.Conn, error) {
    40  	var conn driver.Conn
    41  
    42  	err := backoff.Retry(func() error {
    43  		var err error
    44  		if d.driverName == "postgres" {
    45  			conn, err = pq.DialOpen(keepAliveDialer{}, name)
    46  		} else {
    47  			conn, err = d.Driver.Open(name)
    48  		}
    49  		if err != nil {
    50  			if pqErr, ok := err.(*pq.Error); ok && pqErr.Code.Name() == "too_many_connections" {
    51  				return err
    52  			}
    53  
    54  			return backoff.Permanent(err)
    55  		}
    56  
    57  		return nil
    58  	}, backoff.NewExponentialBackOff())
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  
    63  	return conn, nil
    64  }