github.com/jackc/pgx/v5@v5.5.5/pgconn/pgconn_stress_test.go (about)

     1  package pgconn_test
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"os"
     7  	"runtime"
     8  	"strconv"
     9  	"testing"
    10  
    11  	"github.com/jackc/pgx/v5/pgconn"
    12  
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestConnStress(t *testing.T) {
    17  	pgConn, err := pgconn.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
    18  	require.NoError(t, err)
    19  	defer closeConn(t, pgConn)
    20  
    21  	actionCount := 10000
    22  	if s := os.Getenv("PGX_TEST_STRESS_FACTOR"); s != "" {
    23  		stressFactor, err := strconv.ParseInt(s, 10, 64)
    24  		require.Nil(t, err, "Failed to parse PGX_TEST_STRESS_FACTOR")
    25  		actionCount *= int(stressFactor)
    26  	}
    27  
    28  	setupStressDB(t, pgConn)
    29  
    30  	actions := []struct {
    31  		name string
    32  		fn   func(*pgconn.PgConn) error
    33  	}{
    34  		{"Exec Select", stressExecSelect},
    35  		{"ExecParams Select", stressExecParamsSelect},
    36  		{"Batch", stressBatch},
    37  	}
    38  
    39  	for i := 0; i < actionCount; i++ {
    40  		action := actions[rand.Intn(len(actions))]
    41  		err := action.fn(pgConn)
    42  		require.Nilf(t, err, "%d: %s", i, action.name)
    43  	}
    44  
    45  	// Each call with a context starts a goroutine. Ensure they are cleaned up when context is not canceled.
    46  	numGoroutine := runtime.NumGoroutine()
    47  	require.Truef(t, numGoroutine < 1000, "goroutines appear to be orphaned: %d in process", numGoroutine)
    48  }
    49  
    50  func setupStressDB(t *testing.T, pgConn *pgconn.PgConn) {
    51  	_, err := pgConn.Exec(context.Background(), `
    52  		create temporary table widgets(
    53  			id serial primary key,
    54  			name varchar not null,
    55  			description text,
    56  			creation_time timestamptz default now()
    57  		);
    58  
    59  		insert into widgets(name, description) values
    60  			('Foo', 'bar'),
    61  			('baz', 'Something really long Something really long Something really long Something really long Something really long'),
    62  			('a', 'b')`).ReadAll()
    63  	require.NoError(t, err)
    64  }
    65  
    66  func stressExecSelect(pgConn *pgconn.PgConn) error {
    67  	ctx, cancel := context.WithCancel(context.Background())
    68  	defer cancel()
    69  	_, err := pgConn.Exec(ctx, "select * from widgets").ReadAll()
    70  	return err
    71  }
    72  
    73  func stressExecParamsSelect(pgConn *pgconn.PgConn) error {
    74  	ctx, cancel := context.WithCancel(context.Background())
    75  	defer cancel()
    76  	result := pgConn.ExecParams(ctx, "select * from widgets where id < $1", [][]byte{[]byte("10")}, nil, nil, nil).Read()
    77  	return result.Err
    78  }
    79  
    80  func stressBatch(pgConn *pgconn.PgConn) error {
    81  	ctx, cancel := context.WithCancel(context.Background())
    82  	defer cancel()
    83  
    84  	batch := &pgconn.Batch{}
    85  
    86  	batch.ExecParams("select * from widgets", nil, nil, nil, nil)
    87  	batch.ExecParams("select * from widgets where id < $1", [][]byte{[]byte("10")}, nil, nil, nil)
    88  	_, err := pgConn.ExecBatch(ctx, batch).ReadAll()
    89  	return err
    90  }