Share PGSQL connection pool between ledgers.

To do this, we introduce a storage.Driver interface between the storage.Factory and the Store implementations.
This interface is in charge of initializing the connection to the underlying database (in case of PGSQL, create the pool) and create the underlying stores for ledgers.

ledger/ledger_test.go

@@ -1,9 +1,12 @@
 package ledger
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/jackc/pgx/v4/pgxpool"
+
 	"io/ioutil"
 	"log"
 	"math/rand"
@@ -57,7 +60,14 @@ func TestMain(m *testing.M) {
 		viper.Set("storage.sqlite.db_name", "ledger")
 		os.Remove(path.Join(os.TempDir(), "ledger_test.db"))
 	case "postgres":
-		store, err := postgres.NewStore("test")
+		pool, err := pgxpool.Connect(
+			context.Background(),
+			viper.GetString("storage.postgres.conn_string"),
+		)
+		if err != nil {
+			panic(err)
+		}
+		store, err := postgres.NewStore("test", pool)
 		if err != nil {
 			panic(err)
 		}
@@ -446,7 +456,7 @@ func TestRevertTransaction(t *testing.T) {
 	})
 }
 
-func BenchmarkTransaction1(b *testing.B) {
+func markTransaction1(b *testing.B) {
 	with(func(l *Ledger) {
 		for n := 0; n < b.N; n++ {
 			txs := []core.Transaction{}

storage/postgres/store.go

@@ -9,52 +9,25 @@ import (
 	"strings"
 
 	"github.com/jackc/pgx/v4/pgxpool"
-	"github.com/spf13/viper"
 )
 
 //go:embed migration
 var migrations embed.FS
 
 type PGStore struct {
-	ledger     string
-	connString string
-	pool       *pgxpool.Pool
-}
-
-func (s *PGStore) connect() error {
-	log.Println("initiating postgres pool")
-
-	pool, err := pgxpool.Connect(
-		context.Background(),
-		s.connString,
-	)
-
-	if err != nil {
-		return err
-	}
-
-	s.pool = pool
-
-	return nil
+	ledger string
+	pool   *pgxpool.Pool
 }
 
 func (s *PGStore) Conn() *pgxpool.Pool {
 	return s.pool
 }
 
-func NewStore(name string) (*PGStore, error) {
-	store := &PGStore{
-		ledger:     name,
-		connString: viper.GetString("storage.postgres.conn_string"),
-	}
-
-	err := store.connect()
-
-	if err != nil {
-		return store, err
-	}
-
-	return store, nil
+func NewStore(name string, pool *pgxpool.Pool) (*PGStore, error) {
+	return &PGStore{
+		ledger: name,
+		pool:   pool,
+	}, nil
 }
 
 func (s *PGStore) Initialize() error {
@@ -105,7 +78,7 @@ func (s *PGStore) table(name string) string {
 }
 
 func (s *PGStore) Close() {
-	s.pool.Close()
+	//s.pool.Close()
 }
 
 func (s *PGStore) DropTest() {

storage/storage.go

@@ -1,12 +1,16 @@
 package storage
 
 import (
+	"context"
+	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/numary/ledger/core"
 	"github.com/numary/ledger/ledger/query"
 	"github.com/numary/ledger/storage/postgres"
 	"github.com/numary/ledger/storage/sqlite"
 	"github.com/pkg/errors"
 	"github.com/spf13/viper"
+	"log"
+	"sync"
 )
 
 type Store interface {
@@ -25,22 +29,79 @@ type Store interface {
 	Close()
 }
 
-func GetStore(name string) (Store, error) {
-	driver := viper.GetString("storage.driver")
+type Driver interface {
+	Initialize() error
+	NewStore(name string) (Store, error)
+}
+
+var builtInDrivers = make(map[string]Driver)
+
+func RegisterDriver(name string, driver Driver) {
+	builtInDrivers[name] = driver
+}
 
-	switch driver {
-	case "sqlite":
-		return sqlite.NewStore(name)
-	case "postgres":
-		return postgres.NewStore(name)
+type SQLiteSDriver struct {}
+
+func (d *SQLiteSDriver) Initialize() error {
+	return nil
+}
+
+func (d *SQLiteSDriver) NewStore(name string) (Store, error) {
+	return sqlite.NewStore(name)
+}
+
+type PGSqlDriver struct {
+	once sync.Once
+	pool *pgxpool.Pool
+}
+
+func (d *PGSqlDriver) Initialize() error {
+	errCh := make(chan error, 1)
+	d.once.Do(func() {
+		log.Println("initiating postgres pool")
+
+		pool, err := pgxpool.Connect(
+			context.Background(),
+			viper.GetString("storage.postgres.conn_string"),
+		)
+		if err != nil {
+			errCh <- err
+		}
+		d.pool = pool
+		errCh <- nil
+	})
+	select {
+	case err := <- errCh:
+		return err
 	default:
-		break
+		return nil
+	}
+}
+
+func (d *PGSqlDriver) NewStore(name string) (Store, error) {
+	return postgres.NewStore(name, d.pool)
+}
+
+func init() {
+	RegisterDriver("sqlite", &SQLiteSDriver{})
+	RegisterDriver("postgres", &PGSqlDriver{})
+}
+
+func GetStore(name string) (Store, error) {
+	driverStr := viper.GetString("storage.driver")
+	driver, ok := builtInDrivers[driverStr]
+	if !ok {
+		panic(errors.Errorf(
+			"unsupported store: %s",
+			driver,
+		))
+	}
+	err := driver.Initialize()
+	if err != nil {
+		return nil, errors.Wrap(err, "initializing driver")
 	}
 
-	panic(errors.Errorf(
-		"unsupported store: %s",
-		driver,
-	))
+	return driver.NewStore(name)
 }
 
 type Factory interface {