github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/storage/cloud/gcs_storage_test.go (about)

     1  // Copyright 2020 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  package cloud
    12  
    13  import (
    14  	"context"
    15  	"io/ioutil"
    16  	"math/rand"
    17  	"net"
    18  	"net/http"
    19  	"os"
    20  	"testing"
    21  
    22  	"github.com/cockroachdb/cockroach/pkg/base"
    23  	"github.com/cockroachdb/cockroach/pkg/util/randutil"
    24  	"github.com/cockroachdb/cockroach/pkg/util/sysutil"
    25  	"github.com/stretchr/testify/require"
    26  )
    27  
    28  var econnrefused = &net.OpError{Err: &os.SyscallError{
    29  	Syscall: "test",
    30  	Err:     sysutil.ECONNREFUSED,
    31  }}
    32  
    33  var econnreset = &net.OpError{Err: &os.SyscallError{
    34  	Syscall: "test",
    35  	Err:     sysutil.ECONNRESET,
    36  }}
    37  
    38  type antagonisticDialer struct {
    39  	net.Dialer
    40  	rnd               *rand.Rand
    41  	numRepeatFailures *int
    42  }
    43  
    44  type antagonisticConn struct {
    45  	net.Conn
    46  	rnd               *rand.Rand
    47  	numRepeatFailures *int
    48  }
    49  
    50  func (d *antagonisticDialer) DialContext(
    51  	ctx context.Context, network, addr string,
    52  ) (net.Conn, error) {
    53  	if network == "tcp" && addr == "storage.googleapis.com:443" {
    54  		if *d.numRepeatFailures < maxNoProgressReads && d.rnd.Int()%2 == 0 {
    55  			*(d.numRepeatFailures)++
    56  			return nil, econnrefused
    57  		}
    58  		c, err := d.Dialer.DialContext(ctx, network, addr)
    59  		if err != nil {
    60  			return nil, err
    61  		}
    62  
    63  		return &antagonisticConn{Conn: c, rnd: d.rnd, numRepeatFailures: d.numRepeatFailures}, nil
    64  	}
    65  	return d.Dialer.DialContext(ctx, network, addr)
    66  }
    67  
    68  func (c *antagonisticConn) Read(b []byte) (int, error) {
    69  	if *c.numRepeatFailures < maxNoProgressReads && c.rnd.Int()%2 == 0 {
    70  		*(c.numRepeatFailures)++
    71  		return 0, econnreset
    72  	}
    73  	return c.Conn.Read(b)
    74  }
    75  
    76  func TestAntagonisticRead(t *testing.T) {
    77  	if os.Getenv("GOOGLE_APPLICATION_CREDENTIALS") == "" {
    78  		// This test requires valid GS credential file.
    79  		return
    80  	}
    81  
    82  	rnd, _ := randutil.NewPseudoRand()
    83  	failures := 0
    84  	// Override DialContext implementation in http transport.
    85  	dialer := &antagonisticDialer{rnd: rnd, numRepeatFailures: &failures}
    86  
    87  	// Override transport to return antagonistic connection.
    88  	transport := http.DefaultTransport.(*http.Transport)
    89  	transport.DialContext =
    90  		func(ctx context.Context, network, addr string) (net.Conn, error) {
    91  			return dialer.DialContext(ctx, network, addr)
    92  		}
    93  	transport.DisableKeepAlives = true
    94  
    95  	defer func() {
    96  		transport.DialContext = nil
    97  		transport.DisableKeepAlives = false
    98  	}()
    99  
   100  	gsFile := "gs://cockroach-fixtures/tpch-csv/sf-1/region.tbl?AUTH=implicit"
   101  	conf, err := ExternalStorageConfFromURI(gsFile)
   102  	require.NoError(t, err)
   103  
   104  	s, err := MakeExternalStorage(
   105  		context.Background(), conf, base.ExternalIODirConfig{}, testSettings, nil)
   106  	require.NoError(t, err)
   107  	stream, err := s.ReadFile(context.Background(), "")
   108  	require.NoError(t, err)
   109  	defer stream.Close()
   110  	_, err = ioutil.ReadAll(stream)
   111  	require.NoError(t, err)
   112  }