Merge pull request #5 from iluhinsky/refactoring/db

Moved db-related code to dedicated module

Author: viatoriche

README.md

@@ -4,6 +4,7 @@ perfkit is a set of performance tools.
 
 It includes:
 * [benchmark](benchmark) - a library for writing benchmarks
+* [db](db) - a library for working with databases
 * [acronis-db-bench](acronis-db-bench) - a tool to run database benchmarks
 * [acronis-restrelay-bench](acronis-restrelay-bench) - a tool to run end-to-end cloud deployment benchmarks
 

acronis-db-bench/VERSION

@@ -1 +1 @@
-1.1.1
+1.1.1-beta

acronis-db-bench/acronis-db-bench.go

@@ -9,8 +9,12 @@ import (
 	"sort"
 	"strings"
 
+	_ "net/http/pprof"
+
 	"github.com/acronis/perfkit/benchmark"
-	embeddedpostgres "github.com/fergusstrange/embedded-postgres" // embedder postgres
+	"github.com/acronis/perfkit/db"
+
+	_ "github.com/acronis/perfkit/db/sql" // sql drivers
 )
 
 // Version is a version of the acronis-db-bench
@@ -29,11 +33,10 @@ func printVersion() {
 
 // TestOpts is a structure to store all the test options
 type TestOpts struct {
-	DBOpts               benchmark.DatabaseOpts
-	BenchOpts            BenchOpts
-	EmbeddedPostgresOpts EmbeddedPostgresOpts
-	TestcaseOpts         TestcaseOpts
-	CTIOpts              CTIOpts
+	DBOpts       DatabaseOpts
+	BenchOpts    BenchOpts
+	TestcaseOpts TestcaseOpts
+	CTIOpts      CTIOpts
 }
 
 // BenchOpts is a structure to store all the benchmark options
@@ -70,17 +73,17 @@ type TestcaseOpts struct {
 
 // DBTestData is a structure to store all the test data
 type DBTestData struct {
-	TestDesc         *TestDesc
-	EventBus         *EventBus
-	EmbeddedPostgres *embeddedpostgres.EmbeddedPostgres
-	EffectiveBatch   int // EffectiveBatch reflects the default value if the --batch option is not set, it can be different for different tests
+	TestDesc       *TestDesc
+	EventBus       *EventBus
+	TenantsCache   *TenantsCache
+	EffectiveBatch int // EffectiveBatch reflects the default value if the --batch option is not set, it can be different for different tests
 
 	scores map[string][]benchmark.Score
 }
 
 // DBWorkerData is a structure to store all the worker data
 type DBWorkerData struct {
-	conn *benchmark.DBConnector
+	conn *DBConnector
 }
 
 var header = strings.Repeat("=", 120) + "\n"
@@ -97,17 +100,12 @@ func Main() {
 		var testOpts TestOpts
 		b.Cli.AddFlagGroup("Database options", "", &testOpts.DBOpts)
 		b.Cli.AddFlagGroup("acronis-db-bench specific options", "", &testOpts.BenchOpts)
-		b.Cli.AddFlagGroup("Embedded Postgres specific options", "", &testOpts.EmbeddedPostgresOpts)
 		b.Cli.AddFlagGroup("Testcase specific options", "", &testOpts.TestcaseOpts)
 		b.Cli.AddFlagGroup("CTI-pattern simulation test options", "", &testOpts.CTIOpts)
 
 		return &testOpts
 	}
 
-	b.PreExit = func() {
-		finiEmbeddedPostgres(b)
-	}
-
 	b.PrintScore = func(score benchmark.Score) {
 		testData := b.Vault.(*DBTestData)
 		var format string
@@ -148,15 +146,22 @@ func Main() {
 		b.Vault.(*DBTestData).EffectiveBatch = 1
 	}
 
+	var dialectName, err = db.GetDialectName(b.TestOpts.(*TestOpts).DBOpts.ConnString)
+	if err != nil {
+		b.Exit(err)
+	}
+
 	if testOpts.BenchOpts.List {
 		groups, _ := GetTests()
 		fmt.Printf(header) //nolint:staticcheck
 		for _, g := range groups {
+
 			str := fmt.Sprintf("  -- %s", g.name) //nolint:perfsprint
 			fmt.Printf("\n%s %s\n\n", str, strings.Repeat("-", 130-len(str)))
+
 			var testsOutput []string
 			for _, t := range g.tests {
-				if testOpts.DBOpts.Driver != "" && !t.dbIsSupported(testOpts.DBOpts.Driver) {
+				if dialectName != "" && !t.dbIsSupported(dialectName) {
 					continue
 				}
 				testsOutput = append(testsOutput, fmt.Sprintf("  %-39s : %s : %s\n", t.name, t.getDBs(), t.description))
@@ -166,15 +171,13 @@ func Main() {
 		}
 		fmt.Printf("\n")
 		fmt.Printf("Databases symbol legend:\n\n ")
-		for _, db := range benchmark.GetDatabases() {
+		for _, db := range db.GetDatabases() {
 			fmt.Printf(" %s - %s;", db.Symbol, db.Name)
 		}
 		fmt.Printf("\n\n")
 		b.Exit()
 	}
 
-	initEmbeddedPostgres(b)
-
 	if testOpts.BenchOpts.Describe {
 		describeTest(b, testOpts)
 		b.Exit()
@@ -198,17 +201,24 @@ func Main() {
 	if testOpts.DBOpts.Reconnect {
 		b.PreWorker = func(workerId int) {
 			conn := b.WorkerData[workerId].(*DBWorkerData).conn
-			conn.Close()
+			conn.database.Close()
 		}
 	}
 
 	c := dbConnector(b)
 
-	driver, version := c.GetVersion()
+	driver, version, err := c.database.GetVersion()
+	if err != nil {
+		b.Exit("Failed to get database version: %v", err)
+	}
+
 	fmt.Printf("Connected to '%s' database: %s\n", driver, version)
 	fmt.Printf(header) //nolint:staticcheck
 
-	content, dbInfo := c.GetInfo(version)
+	content, dbInfo, err := c.database.GetInfo(version)
+	if err != nil {
+		b.Exit("Failed to get database info: %v", err)
+	}
 
 	if testOpts.BenchOpts.Info || b.Logger.LogLevel > benchmark.LogInfo {
 		if testOpts.BenchOpts.Info {
@@ -222,7 +232,7 @@ func Main() {
 
 	if testOpts.BenchOpts.ProfilerPort > 0 {
 		go func() {
-			err := http.ListenAndServe(fmt.Sprintf("localhost:%d", testOpts.BenchOpts.ProfilerPort), nil)
+			err = http.ListenAndServe(fmt.Sprintf("localhost:%d", testOpts.BenchOpts.ProfilerPort), nil)
 			if err != nil {
 				b.Exit("Failed to start profiler server: %v", err)
 			}
@@ -232,8 +242,10 @@ func Main() {
 	}
 
 	b.Init = func() {
-		b.TenantsCache.SetTenantsWorkingSet(b.TestOpts.(*TestOpts).BenchOpts.TenantsWorkingSet)
-		b.TenantsCache.SetCTIsWorkingSet(b.TestOpts.(*TestOpts).BenchOpts.CTIsWorkingSet)
+		b.Vault.(*DBTestData).TenantsCache = NewTenantsCache(b)
+
+		b.Vault.(*DBTestData).TenantsCache.SetTenantsWorkingSet(b.TestOpts.(*TestOpts).BenchOpts.TenantsWorkingSet)
+		b.Vault.(*DBTestData).TenantsCache.SetCTIsWorkingSet(b.TestOpts.(*TestOpts).BenchOpts.CTIsWorkingSet)
 
 		if b.Logger.LogLevel > benchmark.LogInfo && !testOpts.BenchOpts.Info {
 			b.Log(benchmark.LogTrace, 0, getDBInfo(b, content))
@@ -298,9 +310,15 @@ func executeTests(b *benchmark.Benchmark, testOpts *TestOpts) {
 	if !exists {
 		b.Exit(fmt.Sprintf("Test: '%s' doesn't exist, see the list of available tests using --list option\n", testOpts.BenchOpts.Test))
 	}
+
+	var dialectName, err = db.GetDialectName(testOpts.DBOpts.ConnString)
+	if err != nil {
+		b.Exit(err)
+	}
+
 	test := tests[testOpts.BenchOpts.Test]
-	if !test.dbIsSupported(testOpts.DBOpts.Driver) {
-		b.Exit(fmt.Sprintf("Test: '%s' doesn't support '%s' database\n", testOpts.BenchOpts.Test, testOpts.DBOpts.Driver))
+	if !test.dbIsSupported(dialectName) {
+		b.Exit(fmt.Sprintf("Test: '%s' doesn't support '%s' database\n", testOpts.BenchOpts.Test, dialectName))
 	}
 	test.launcherFunc(b, test)
 }

acronis-db-bench/db.go

@@ -0,0 +1,204 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/acronis/perfkit/benchmark"
+	"github.com/acronis/perfkit/db"
+)
+
+const SequenceName = "acronis_db_bench_sequence" // SequenceName is the name of the sequence used for generating IDs
+
+// DatabaseOpts represents common flags for every test
+type DatabaseOpts struct {
+	ConnString   string `long:"connection-string" description:"connection string" default:"sqlite://:memory:" required:"false"`
+	MaxOpenConns int    `long:"max-open-cons" description:"max open connections per worker" default:"2" required:"false"`
+	Reconnect    bool   `long:"reconnect" description:"reconnect to DB before every test iteration" required:"false"`
+
+	DryRun bool `long:"dry-run" description:"do not execute any INSERT/UPDATE/DELETE queries on DB-side" required:"false"`
+
+	LogQueries    bool `long:"log-queries" description:"log queries" required:"false"`
+	LogReadedRows bool `long:"log-readed-rows" description:"log readed rows" required:"false"`
+	LogQueryTime  bool `long:"log-query-time" description:"log query time" required:"false"`
+
+	DontCleanup bool `long:"dont-cleanup" description:"do not cleanup DB content before/after the test in '-t all' mode" required:"false"`
+	UseTruncate bool `long:"use-truncate" description:"use TRUNCATE instead of DROP TABLE in cleanup procedure" required:"false"`
+}
+
+// dbConnectorsPool is a simple connection pool, required not to saturate DB connection pool
+type dbConnectorsPool struct {
+	lock sync.Mutex
+	pool map[string]*DBConnector
+}
+
+// key returns a unique key for the connection pool
+func (p *dbConnectorsPool) key(connectionString string, workerID int) string {
+	return fmt.Sprintf("%s-%d", connectionString, workerID)
+}
+
+// take returns a connection from the pool or nil if the pool is empty
+func (p *dbConnectorsPool) take(dbOpts *DatabaseOpts, workerID int) *DBConnector {
+	k := p.key(dbOpts.ConnString, workerID)
+
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	conn, exists := p.pool[k]
+	if exists {
+		delete(p.pool, k)
+		conn.Log(benchmark.LogTrace, "taking connection from the connection pool")
+
+		return conn
+	}
+
+	return nil
+}
+
+// put puts a connection to the pool
+func (p *dbConnectorsPool) put(conn *DBConnector) {
+	k := p.key(conn.DbOpts.ConnString, conn.WorkerID)
+
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	_, exists := p.pool[k]
+	if exists {
+		FatalError("trying to put connection while another connection in the pool already exists")
+	}
+	p.pool[k] = conn
+	conn.Log(benchmark.LogTrace, "releasing connection to the connection pool")
+}
+
+// newDBConnectorsPool creates a new connection pool
+func newDBConnectorsPool() *dbConnectorsPool {
+	p := dbConnectorsPool{}
+	p.pool = make(map[string]*DBConnector)
+
+	return &p
+}
+
+// connPool is a global connection pool
+var connPool = newDBConnectorsPool()
+
+// connectionsChecker checks for potential connections leak
+func connectionsChecker(conn *DBConnector) {
+	for {
+		if conn.database != nil {
+			openConnections := 0
+
+			conn.lock.Lock()
+			if conn.database != nil {
+				stats := conn.database.Stats()
+				openConnections = stats.OpenConnections
+			}
+			conn.lock.Unlock()
+
+			if openConnections > 1 {
+				conn.Log(benchmark.LogError, "internal error: potential connections leak detected")
+			}
+		}
+		time.Sleep(3 * time.Second)
+	}
+}
+
+/*
+ * DB connection management
+ */
+
+// DBConnector is a wrapper for DB connection
+type DBConnector struct {
+	Logger        *benchmark.Logger
+	DbOpts        *DatabaseOpts
+	RetryAttempts int
+	WorkerID      int
+
+	lock     sync.Mutex
+	database db.Database
+}
+
+// NewDBConnector creates a new DBConnector
+func NewDBConnector(dbOpts *DatabaseOpts, workerID int, logger *benchmark.Logger, retryAttempts int) (*DBConnector, error) {
+	c := connPool.take(dbOpts, workerID)
+	if c != nil {
+		return c, nil
+	}
+
+	var queryLogger, readedRowsLogger, queryTimeLogger db.Logger
+	if dbOpts.LogQueries {
+		queryLogger = &dbLogger{level: benchmark.LogTrace, worker: workerID, logger: logger}
+	}
+
+	if dbOpts.LogReadedRows {
+		readedRowsLogger = &dbLogger{level: benchmark.LogTrace, worker: workerID, logger: logger}
+	}
+
+	if dbOpts.LogQueryTime {
+		queryTimeLogger = &dbLogger{level: benchmark.LogTrace, worker: workerID, logger: logger}
+	}
+
+	var dbConn, err = db.Open(db.Config{
+		ConnString:   dbOpts.ConnString,
+		MaxOpenConns: dbOpts.MaxOpenConns,
+		DryRun:       dbOpts.DryRun,
+		UseTruncate:  dbOpts.UseTruncate,
+
+		QueryLogger:      queryLogger,
+		ReadedRowsLogger: readedRowsLogger,
+		QueryTimeLogger:  queryTimeLogger,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	c = &DBConnector{
+		Logger:        logger,
+		DbOpts:        dbOpts,
+		RetryAttempts: retryAttempts,
+		WorkerID:      workerID,
+		database:      dbConn,
+	}
+
+	// go connectionsChecker(c)
+
+	return c, nil
+}
+
+// Release releases the connection to the pool
+func (c *DBConnector) Release() {
+	connPool.put(c)
+}
+
+// SetLogLevel sets log level
+func (c *DBConnector) SetLogLevel(logLevel int) {
+	c.Logger.LogLevel = logLevel
+}
+
+// Log logs a message
+func (c *DBConnector) Log(LogLevel int, format string, args ...interface{}) {
+	c.Logger.Log(LogLevel, c.WorkerID, format, args...)
+}
+
+// Logn logs a message without a newline
+func (c *DBConnector) Logn(LogLevel int, format string, args ...interface{}) {
+	c.Logger.Logn(LogLevel, c.WorkerID, format, args...)
+}
+
+// Exit exits with an error message
+func (c *DBConnector) Exit(fmts string, args ...interface{}) {
+	fmt.Printf(fmts, args...)
+	fmt.Println()
+	os.Exit(127)
+}
+
+type dbLogger struct {
+	level  int
+	worker int
+	logger *benchmark.Logger
+}
+
+func (l *dbLogger) Log(format string, args ...interface{}) {
+	l.logger.Log(l.level, l.worker, format, args...)
+}