allow to use pool of ydb servers

Author: rekby

tools/greenplum-to-pg-tests/cmd/check_pg_queries.go

@@ -56,7 +56,7 @@ func init() {
 	checkPgQueriesCmd.PersistentFlags().BoolVar(&checkPgQueriesConfig.sessionsLogNeedSort, "query-log-need-sort", false, "Sort query log in memory before start")
 	must0(checkPgQueriesCmd.MarkPersistentFlagRequired("query-log"))
 
-	checkPgQueriesCmd.PersistentFlags().BoolVar(&checkPgQueriesConfig.includeFailed, "include-failed", false, "Extract sessions with failed transactions")
+	checkPgQueriesCmd.PersistentFlags().BoolVar(&checkPgQueriesConfig.includeFailed, "include-failed", true, "Extract sessions with failed transactions")
 	checkPgQueriesCmd.PersistentFlags().StringVar(&checkPgQueriesConfig.ydbConnectionString, "ydb-connection", "grpc://localhost:2136/local", "Connection string to ydb server for check queries")
 	checkPgQueriesCmd.PersistentFlags().IntVar(&checkPgQueriesConfig.limitRequests, "requests-limit", 0, "Limit number of parse requests, 0 mean unlimited")
 	checkPgQueriesCmd.PersistentFlags().StringVar(&checkPgQueriesConfig.rulesFile, "rules-file", "issues.yaml", "Rules for detect issue. Set empty for skip read rules.")
@@ -106,10 +106,8 @@ var checkPgQueriesCmd = &cobra.Command{
 
 		log.Println("Connecting to ydb...")
 		connectCtx, cancel := context.WithTimeout(ctx, time.Second*10)
-		db := must(ydb.Open(
-			connectCtx, checkPgQueriesConfig.ydbConnectionString,
-			internal.GetYdbCredentials(),
-		))
+		connectionStrings := strings.Split(checkPgQueriesConfig.ydbConnectionString, ",")
+		dbPool := internal.OpenYdbPool(connectCtx, connectionStrings, []ydb.Option{internal.GetYdbCredentials()})
 		cancel()
 
 		var queries <-chan string
@@ -122,7 +120,7 @@ var checkPgQueriesCmd = &cobra.Command{
 
 		log.Println("Start check queries")
 		var stats QueryStats
-		checkQueries(rules, &stats, db, queries)
+		checkQueries(rules, &stats, dbPool, queries)
 
 		if checkPgQueriesConfig.writeRulesWithStat != "" {
 			rules.UpdateFromStats(stats, checkPgQueriesConfig.sortRulesByCount)
@@ -210,6 +208,9 @@ func readSortedQueries(reader io.ReadCloser) <-chan string {
 					// pass
 				}
 			}
+			if !item.TransactionSuccess && !checkPgQueriesConfig.includeFailed {
+				continue
+			}
 
 			queries <- item.Query
 			counter++
@@ -345,6 +346,10 @@ func extractQueries(sessions []internal.Session) <-chan string {
 		needRemoveLine := false
 		for _, session := range sessions {
 			for _, transaction := range session.Transactions {
+				if !transaction.Success && !checkPgQueriesConfig.includeFailed {
+					continue
+				}
+
 				for _, pgQuery := range transaction.Queries {
 					queryIndex++
 					if queryIndex%checkPgQueriesConfig.printProgressEveryQueries == 0 {
@@ -371,7 +376,7 @@ func printDeleteLine() {
 	fmt.Printf("\033[1A\033[K")
 }
 
-func checkQueries(rules Rules, stats *QueryStats, db *ydb.Driver, queries <-chan string) {
+func checkQueries(rules Rules, stats *QueryStats, dbPool *internal.YdbPool, queries <-chan string) {
 	if checkPgQueriesConfig.checkersCount < 1 {
 		log.Fatalf("can't start less then 1 checker, got: %v", checkPgQueriesConfig.checkersCount)
 	}
@@ -384,7 +389,7 @@ func checkQueries(rules Rules, stats *QueryStats, db *ydb.Driver, queries <-chan
 		go func() {
 			defer wg.Done()
 			for q := range queries {
-				checkQuery(stats, rules, db, q)
+				checkQuery(stats, rules, dbPool, q)
 				counter := itemsCounter.Add(1)
 				if counter%writeStatEveryItems == 0 && checkPgQueriesConfig.writeStatPath != "" {
 					if err := stats.SaveToFile(checkPgQueriesConfig.writeStatPath); err != nil {
@@ -405,7 +410,10 @@ const (
 	checkResultErrUnknown
 )
 
-func checkQuery(stat *QueryStats, rules Rules, db *ydb.Driver, queryText string) (reason string, checkResult checkResultType) {
+func checkQuery(stat *QueryStats, rules Rules, dbPool *internal.YdbPool, queryText string) (reason string, checkResult checkResultType) {
+	db := dbPool.Get()
+	defer dbPool.Release(db)
+
 	queryText = strings.TrimSpace(queryText)
 	queryText = fixSchemaNames(queryText)
 	queryText = fixCreateTable(queryText)

tools/greenplum-to-pg-tests/internal/ydb_pool.go

@@ -0,0 +1,98 @@
+package internal
+
+import (
+	"context"
+	"log"
+	"sync"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3"
+)
+
+type YdbPool struct {
+	m       sync.Mutex
+	drivers []driverInflight
+}
+
+type driverInflight struct {
+	driver   *ydb.Driver
+	inflyght int
+}
+
+func OpenYdbPool(ctx context.Context, connectionStrings []string, options []ydb.Option) *YdbPool {
+	var m sync.Mutex
+	drivers := make([]*ydb.Driver, 0, len(connectionStrings))
+
+	var wg sync.WaitGroup
+	for i := range connectionStrings {
+		cs := connectionStrings[i]
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			driver, err := ydb.Open(ctx, cs, options...)
+			if err == nil {
+				m.Lock()
+				drivers = append(drivers, driver)
+				m.Unlock()
+			} else {
+				log.Printf("Failed connection to %q: %v", cs, err)
+			}
+		}()
+	}
+	wg.Wait()
+
+	if len(drivers) == 0 {
+		log.Fatalf("failed to connect for all endpoints")
+	}
+
+	return NewYdbPool(drivers)
+}
+
+func NewYdbPool(drivers []*ydb.Driver) *YdbPool {
+	if len(drivers) == 0 {
+		panic("can't create pool without drivers")
+	}
+	res := &YdbPool{
+		drivers: make([]driverInflight, len(drivers)),
+	}
+	for i, driver := range drivers {
+		res.drivers[i] = driverInflight{
+			driver:   driver,
+			inflyght: 0,
+		}
+	}
+	return res
+}
+
+// Get return driver with minimal current inflight
+func (p *YdbPool) Get() *ydb.Driver {
+	p.m.Lock()
+	defer p.m.Unlock()
+
+	minInflyght := &p.drivers[0]
+	for i := range p.drivers {
+		if p.drivers[i].inflyght < minInflyght.inflyght {
+			minInflyght = &p.drivers[i]
+		}
+	}
+
+	minInflyght.inflyght++
+	return minInflyght.driver
+}
+
+func (p *YdbPool) Release(driver *ydb.Driver) {
+	p.m.Lock()
+	defer p.m.Unlock()
+
+	for i := range p.drivers {
+		if p.drivers[i].driver == driver {
+			p.drivers[i].inflyght--
+			if p.drivers[i].inflyght < 0 {
+				panic("the ydb driver released more times, then get")
+			}
+			return
+		}
+	}
+
+	panic("the driver doesn't exists in the pool")
+}