github.com/blend/go-sdk@v1.20220411.3/examples/db/prevent-deadlock/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  	"database/sql"
    13  	"log"
    14  	"sync"
    15  	"time"
    16  
    17  	"github.com/blend/go-sdk/db"
    18  	"github.com/blend/go-sdk/ex"
    19  )
    20  
    21  const (
    22  	updateRows = "UPDATE might_deadlock SET counter = counter + 1 WHERE key = $1;"
    23  	separator  = "=================================================="
    24  )
    25  
    26  func createConn(ctx context.Context) (*db.Connection, error) {
    27  	pool, err := db.New(db.OptConfigFromEnv())
    28  	if err != nil {
    29  		return nil, err
    30  	}
    31  
    32  	err = pool.Open()
    33  	if err != nil {
    34  		return nil, err
    35  	}
    36  
    37  	err = pool.Connection.PingContext(ctx)
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  
    42  	log.Printf("DSN=%q\n", pool.Config.CreateDSN())
    43  	return pool, nil
    44  }
    45  
    46  // contendReads introduces two reads (in a transaction) with a sleep in
    47  // between.
    48  // H/T to https://www.citusdata.com/blog/2018/02/22/seven-tips-for-dealing-with-postgres-locks/
    49  // for the idea on how to "easily" introduce a deadlock.
    50  func contendReads(ctx context.Context, wg *sync.WaitGroup, tx *sql.Tx, key1, key2 string, cfg *config) error {
    51  	defer wg.Done()
    52  
    53  	_, err := tx.ExecContext(ctx, updateRows, key1)
    54  	if err != nil {
    55  		return err
    56  	}
    57  
    58  	time.Sleep(cfg.TxSleep)
    59  	_, err = tx.ExecContext(ctx, updateRows, key2)
    60  	if err == context.DeadlineExceeded {
    61  		return nest(err, ex.New("Context cancel in between queries"))
    62  	}
    63  	return err
    64  }
    65  
    66  func intentionalContention(ctx context.Context, pool *db.Connection, cfg *config) (err error) {
    67  	var tx1, tx2 *sql.Tx
    68  	defer func() {
    69  		err = txFinalize(tx1, err)
    70  		err = txFinalize(tx2, err)
    71  	}()
    72  
    73  	log.Println("Starting transactions")
    74  	tx1, err = pool.BeginContext(ctx)
    75  	if err != nil {
    76  		return
    77  	}
    78  	tx2, err = pool.BeginContext(ctx)
    79  	if err != nil {
    80  		return
    81  	}
    82  	log.Println("Transactions opened")
    83  
    84  	// Kick off two goroutines that contend with each other.
    85  	wg := sync.WaitGroup{}
    86  	errLock := sync.Mutex{}
    87  	wg.Add(2)
    88  	go func() {
    89  		contendErr := contendReads(ctx, &wg, tx1, "hello", "world", cfg)
    90  		errLock.Lock()
    91  		defer errLock.Unlock()
    92  		err = nest(err, contendErr)
    93  	}()
    94  	go func() {
    95  		contendErr := contendReads(ctx, &wg, tx2, "world", "hello", cfg)
    96  		errLock.Lock()
    97  		defer errLock.Unlock()
    98  		err = nest(err, contendErr)
    99  	}()
   100  	wg.Wait()
   101  
   102  	if err != nil {
   103  		return
   104  	}
   105  
   106  	// Make sure to commit both transactions before moving on.
   107  	err = nest(err, tx1.Commit())
   108  	err = nest(err, tx2.Commit())
   109  	return
   110  }
   111  
   112  func main() {
   113  	log.SetFlags(0)
   114  	log.SetOutput(newLogWriter())
   115  	cfg := getConfig()
   116  
   117  	// 1. Set the `DB_LOCK_TIMEOUT` environment variable.
   118  	log.Println(separator)
   119  	cfg.Print()
   120  	err := cfg.SetEnvironment()
   121  	if err != nil {
   122  		log.Fatal(err)
   123  	}
   124  
   125  	deadline := time.Now().Add(cfg.ContextTimeout)
   126  	ctx, cancel := context.WithDeadline(context.Background(), deadline)
   127  	defer cancel()
   128  
   129  	// 2. Parse config / open / ping
   130  	// 3. Make sure `lock_timeout` is in the connection string (it gets printed)
   131  	log.Println(separator)
   132  	pool, err := createConn(ctx)
   133  	if err != nil {
   134  		log.Fatal(err)
   135  	}
   136  	defer cleanUp(pool)
   137  
   138  	// 4. Demonstrate that the observed lock timeout on an open connection is
   139  	//    `LockTimeout`.
   140  	log.Println(separator)
   141  	timeout, err := ensureLockTimeout(ctx, pool, cfg)
   142  	if err != nil {
   143  		log.Fatal(err)
   144  	}
   145  	log.Printf("lock_timeout=%s\n", timeout)
   146  
   147  	// 5. Create a table schema and insert data to seed the database.
   148  	err = seedDatabase(ctx, pool)
   149  	if err != nil {
   150  		log.Fatal(err)
   151  	}
   152  
   153  	// 6. Create two goroutines that intentionally contend with transactions.
   154  	log.Println(separator)
   155  	err = intentionalContention(ctx, pool, cfg)
   156  	if err == nil {
   157  		log.Fatal(ex.New("Expected lock contention to occur"))
   158  	}
   159  
   160  	// 7. Display the error / errors in as verbose a way as possible.
   161  	log.Println("***")
   162  	err = displayError(err)
   163  	if err != nil {
   164  		log.Fatal(err)
   165  	}
   166  }