github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/utils/concurrency/runner.go (about)

     1  // Copyright 2022 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package main
    16  
    17  import (
    18  	"context"
    19  	"database/sql"
    20  	"fmt"
    21  	"os"
    22  
    23  	_ "github.com/go-sql-driver/mysql"
    24  	"golang.org/x/sync/errgroup"
    25  )
    26  
    27  const clients = 16
    28  const iters = 10
    29  
    30  var sqlScript = []string{
    31  	"call dolt_checkout('main');",
    32  	"select * from dolt_log order by date desc limit 10;",
    33  }
    34  
    35  var (
    36  	database = "SHAQ"
    37  	user     = "root"
    38  	pass     = ""
    39  	host     = "127.0.0.1"
    40  	port     = "3306"
    41  )
    42  
    43  // Runs |sqlScript| concurrently on multiple clients.
    44  // Useful for repoducing concurrency bugs.
    45  func main() {
    46  	connStr := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s",
    47  		user, pass, host, port, database)
    48  
    49  	db, err := sql.Open("mysql", connStr)
    50  	maybeExit(err)
    51  
    52  	eg, ctx := errgroup.WithContext(context.Background())
    53  
    54  	for i := 0; i < clients; i++ {
    55  		eg.Go(func() (err error) {
    56  			conn, err := db.Conn(ctx)
    57  			if err != nil {
    58  				return err
    59  			}
    60  			defer func() {
    61  				cerr := conn.Close()
    62  				if err != nil {
    63  					err = cerr
    64  				}
    65  			}()
    66  			for j := 0; j < iters; j++ {
    67  				if err = query(ctx, conn); err != nil {
    68  					return err
    69  				}
    70  			}
    71  			return
    72  		})
    73  	}
    74  	maybeExit(eg.Wait())
    75  }
    76  
    77  func query(ctx context.Context, conn *sql.Conn) error {
    78  	for i := range sqlScript {
    79  		_, err := conn.ExecContext(ctx, sqlScript[i])
    80  		if err != nil {
    81  			return err
    82  		}
    83  	}
    84  	return nil
    85  }
    86  
    87  func maybeExit(err error) {
    88  	if err != nil {
    89  		fmt.Println(err.Error())
    90  		os.Exit(1)
    91  	}
    92  }