github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/pgwire_internal_test.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  // This file contains tests for pgwire that need to be in the sql package.
    12  
    13  package sql
    14  
    15  import (
    16  	"context"
    17  	"database/sql/driver"
    18  	"net/url"
    19  	"testing"
    20  
    21  	"github.com/cockroachdb/cockroach/pkg/base"
    22  	"github.com/cockroachdb/cockroach/pkg/keys"
    23  	"github.com/cockroachdb/cockroach/pkg/security"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/catalog/lease"
    25  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    26  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    27  	"github.com/cockroachdb/cockroach/pkg/testutils"
    28  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    29  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    30  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    31  	"github.com/cockroachdb/errors"
    32  	"github.com/lib/pq"
    33  )
    34  
    35  // Test that abruptly closing a pgwire connection releases all leases held by
    36  // that session.
    37  func TestPGWireConnectionCloseReleasesLeases(t *testing.T) {
    38  	defer leaktest.AfterTest(t)()
    39  	s, _, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
    40  	ctx := context.Background()
    41  	defer s.Stopper().Stop(ctx)
    42  	url, cleanupConn := sqlutils.PGUrl(t, s.ServingSQLAddr(), "SetupServer", url.User(security.RootUser))
    43  	defer cleanupConn()
    44  	conn, err := pq.Open(url.String())
    45  	if err != nil {
    46  		t.Fatal(err)
    47  	}
    48  	ex := conn.(driver.ExecerContext)
    49  	if _, err := ex.ExecContext(ctx, "CREATE DATABASE test", nil); err != nil {
    50  		t.Fatal(err)
    51  	}
    52  	if _, err := ex.ExecContext(ctx, "CREATE TABLE test.t (i INT PRIMARY KEY)", nil); err != nil {
    53  		t.Fatal(err)
    54  	}
    55  	// Start a txn so leases are accumulated by queries.
    56  	if _, err := ex.ExecContext(ctx, "BEGIN", nil); err != nil {
    57  		t.Fatal(err)
    58  	}
    59  	// Get a table lease.
    60  	if _, err := ex.ExecContext(ctx, "SELECT * FROM test.t", nil); err != nil {
    61  		t.Fatal(err)
    62  	}
    63  	// Abruptly close the connection.
    64  	if err := conn.Close(); err != nil {
    65  		t.Fatal(err)
    66  	}
    67  	// Verify that there are no leases held.
    68  	tableDesc := sqlbase.GetTableDescriptor(kvDB, keys.SystemSQLCodec, "test", "t")
    69  
    70  	lm := s.LeaseManager().(*lease.Manager)
    71  
    72  	// Looking for a table state validates that there used to be a lease on the
    73  	// table.
    74  	var leases int
    75  	lm.VisitLeases(func(
    76  		desc sqlbase.TableDescriptor, dropped bool, refCount int, expiration tree.DTimestamp,
    77  	) (wantMore bool) {
    78  		if desc.ID == tableDesc.ID {
    79  			leases++
    80  		}
    81  		return true
    82  	})
    83  	if leases != 1 {
    84  		t.Fatalf("expected one lease, found: %d", leases)
    85  	}
    86  
    87  	// Wait for the lease to be released.
    88  	testutils.SucceedsSoon(t, func() error {
    89  		var totalRefCount int
    90  		lm.VisitLeases(func(
    91  			desc sqlbase.TableDescriptor, dropped bool, refCount int, expiration tree.DTimestamp,
    92  		) (wantMore bool) {
    93  			if desc.ID == tableDesc.ID {
    94  				totalRefCount += refCount
    95  			}
    96  			return true
    97  		})
    98  		if totalRefCount != 0 {
    99  			return errors.Errorf(
   100  				"expected lease to be unused, found refcount: %d", totalRefCount)
   101  		}
   102  		return nil
   103  	})
   104  }