github.com/blend/go-sdk@v1.20220411.3/examples/db/bench/main.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package main
     9  
    10  import (
    11  	"context"
    12  	"fmt"
    13  	"os"
    14  	"strconv"
    15  	"strings"
    16  	"sync"
    17  	"time"
    18  
    19  	"github.com/blend/go-sdk/db"
    20  	"github.com/blend/go-sdk/db/migration"
    21  	"github.com/blend/go-sdk/logger"
    22  	"github.com/blend/go-sdk/stringutil"
    23  )
    24  
    25  func createTable(tableName string, log logger.Log, conn *db.Connection) error {
    26  	log.Infof("creating %s", tableName)
    27  	action := migration.NewStep(
    28  		migration.TableNotExists(tableName),
    29  		migration.Statements(
    30  			fmt.Sprintf("CREATE TABLE %s (id serial not null primary key, name varchar(255))", tableName)),
    31  	)
    32  
    33  	suite := migration.New(migration.OptGroups(migration.NewGroup(migration.OptGroupActions(action))))
    34  	suite.Log = log
    35  	return suite.Apply(context.TODO(), conn)
    36  }
    37  
    38  func dropTable(tableName string, log logger.Log, conn *db.Connection) error {
    39  	log.Infof("dropping %s", tableName)
    40  	action := migration.NewStep(
    41  		migration.TableExists(tableName),
    42  		migration.Statements(
    43  			fmt.Sprintf("DROP TABLE %s", tableName),
    44  		),
    45  	)
    46  
    47  	suite := migration.New(migration.OptGroups(migration.NewGroup(migration.OptGroupActions(action))))
    48  	suite.Log = log
    49  	return suite.Apply(context.TODO(), conn)
    50  }
    51  
    52  func maybeFatal(err error) {
    53  	if err != nil {
    54  		fmt.Fprintf(os.Stderr, "%+v\n", err)
    55  		os.Exit(1)
    56  	}
    57  }
    58  
    59  func reportStats(log logger.Log, conn *db.Connection) {
    60  	ticker := time.Tick(500 * time.Millisecond)
    61  	for range ticker {
    62  		log.Infof("[%d] connections currently open", conn.Connection.Stats().OpenConnections)
    63  		log.Infof("[%v] wait duration", conn.Connection.Stats().WaitDuration)
    64  	}
    65  }
    66  
    67  func main() {
    68  	log := logger.All(logger.OptDisabled(db.QueryFlag))
    69  	conn, err := db.New(db.OptConfigFromEnv())
    70  	if err != nil {
    71  		log.Fatal(err)
    72  		os.Exit(1)
    73  	}
    74  
    75  	if err := conn.Open(); err != nil {
    76  		log.Fatal(err)
    77  		os.Exit(1)
    78  	}
    79  
    80  	go reportStats(log, conn)
    81  
    82  	tableName := strings.ToLower(stringutil.Random(stringutil.Letters, 12))
    83  
    84  	maybeFatal(createTable(tableName, log, conn))
    85  	defer func() { maybeFatal(dropTable(tableName, log, conn)) }()
    86  
    87  	for x := 0; x < 1<<12; x++ {
    88  		ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond)
    89  		err := db.IgnoreExecResult(conn.Invoke(db.OptContext(ctx)).Exec(fmt.Sprintf("INSERT INTO %s VALUES ($1)", tableName), strconv.Itoa(x)))
    90  		maybeFatal(err)
    91  		cancel()
    92  	}
    93  
    94  	wg := sync.WaitGroup{}
    95  	wg.Add(4)
    96  	for routine := 0; routine < 4; routine++ {
    97  		go func() {
    98  			defer wg.Done()
    99  			for x := 0; x < 1<<10; x++ {
   100  				ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond)
   101  				if _, err := conn.Invoke(db.OptContext(ctx)).Query(fmt.Sprintf("select * from %s", tableName)).Any(); err != nil {
   102  					maybeFatal(err)
   103  				}
   104  				cancel()
   105  			}
   106  		}()
   107  	}
   108  	wg.Wait()
   109  
   110  	log.DrainContext(context.TODO())
   111  	log.Infof("OK")
   112  }