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

     1  // Copyright 2019 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 blobs
    12  
    13  import (
    14  	"bytes"
    15  	"context"
    16  	"io/ioutil"
    17  	"os"
    18  	"path/filepath"
    19  	"testing"
    20  	"time"
    21  
    22  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    23  	"github.com/cockroachdb/cockroach/pkg/rpc"
    24  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    25  )
    26  
    27  // filesize should be at least 1 GB when running these benchmarks.
    28  // Reduced to 129 K for CI.
    29  const filesize = 129 * 1 << 10
    30  
    31  type benchmarkTestCase struct {
    32  	localNodeID       roachpb.NodeID
    33  	remoteNodeID      roachpb.NodeID
    34  	localExternalDir  string
    35  	remoteExternalDir string
    36  
    37  	blobClient BlobClient
    38  	fileSize   int64
    39  	fileName   string
    40  }
    41  
    42  func writeLargeFile(t testing.TB, file string, size int64) {
    43  	err := os.MkdirAll(filepath.Dir(file), 0755)
    44  	if err != nil {
    45  		t.Fatal(err)
    46  	}
    47  	content := make([]byte, size)
    48  	err = ioutil.WriteFile(file, content, 0600)
    49  	if err != nil {
    50  		t.Fatal(err)
    51  	}
    52  }
    53  
    54  func BenchmarkStreamingReadFile(b *testing.B) {
    55  	localNodeID := roachpb.NodeID(1)
    56  	remoteNodeID := roachpb.NodeID(2)
    57  	localExternalDir, remoteExternalDir, stopper, cleanUpFn := createTestResources(b)
    58  	defer cleanUpFn()
    59  
    60  	clock := hlc.NewClock(hlc.UnixNano, time.Nanosecond)
    61  	rpcContext := rpc.NewInsecureTestingContext(clock, stopper)
    62  	rpcContext.TestingAllowNamedRPCToAnonymousServer = true
    63  
    64  	factory := setUpService(b, rpcContext, localNodeID, remoteNodeID, localExternalDir, remoteExternalDir)
    65  	blobClient, err := factory(context.Background(), remoteNodeID)
    66  	if err != nil {
    67  		b.Fatal(err)
    68  	}
    69  	params := &benchmarkTestCase{
    70  		localNodeID:       localNodeID,
    71  		remoteNodeID:      remoteNodeID,
    72  		localExternalDir:  localExternalDir,
    73  		remoteExternalDir: remoteExternalDir,
    74  		blobClient:        blobClient,
    75  		fileSize:          filesize,
    76  		fileName:          "test/largefile.csv",
    77  	}
    78  	benchmarkStreamingReadFile(b, params)
    79  }
    80  
    81  func benchmarkStreamingReadFile(b *testing.B, tc *benchmarkTestCase) {
    82  	writeLargeFile(b, filepath.Join(tc.remoteExternalDir, tc.fileName), tc.fileSize)
    83  	writeTo := LocalStorage{externalIODir: tc.localExternalDir}
    84  	b.ResetTimer()
    85  	b.SetBytes(tc.fileSize)
    86  	for i := 0; i < b.N; i++ {
    87  		reader, err := tc.blobClient.ReadFile(context.Background(), tc.fileName)
    88  		if err != nil {
    89  			b.Fatal(err)
    90  		}
    91  		err = writeTo.WriteFile(tc.fileName, reader)
    92  		if err != nil {
    93  			b.Fatal(err)
    94  		}
    95  		stat, err := writeTo.Stat(tc.fileName)
    96  		if err != nil {
    97  			b.Fatal(err)
    98  		}
    99  		if stat.Filesize != tc.fileSize {
   100  			b.Fatal("incorrect number of bytes written")
   101  		}
   102  	}
   103  }
   104  
   105  func BenchmarkStreamingWriteFile(b *testing.B) {
   106  	localNodeID := roachpb.NodeID(1)
   107  	remoteNodeID := roachpb.NodeID(2)
   108  	localExternalDir, remoteExternalDir, stopper, cleanUpFn := createTestResources(b)
   109  	defer cleanUpFn()
   110  
   111  	clock := hlc.NewClock(hlc.UnixNano, time.Nanosecond)
   112  	rpcContext := rpc.NewInsecureTestingContext(clock, stopper)
   113  	rpcContext.TestingAllowNamedRPCToAnonymousServer = true
   114  
   115  	factory := setUpService(b, rpcContext, localNodeID, remoteNodeID, localExternalDir, remoteExternalDir)
   116  	blobClient, err := factory(context.Background(), remoteNodeID)
   117  	if err != nil {
   118  		b.Fatal(err)
   119  	}
   120  	params := &benchmarkTestCase{
   121  		localNodeID:       localNodeID,
   122  		remoteNodeID:      remoteNodeID,
   123  		localExternalDir:  localExternalDir,
   124  		remoteExternalDir: remoteExternalDir,
   125  		blobClient:        blobClient,
   126  		fileSize:          filesize,
   127  		fileName:          "test/largefile.csv",
   128  	}
   129  	benchmarkStreamingWriteFile(b, params)
   130  }
   131  
   132  func benchmarkStreamingWriteFile(b *testing.B, tc *benchmarkTestCase) {
   133  	content := make([]byte, tc.fileSize)
   134  	b.ResetTimer()
   135  	b.SetBytes(tc.fileSize)
   136  	for i := 0; i < b.N; i++ {
   137  		err := tc.blobClient.WriteFile(context.Background(), tc.fileName, bytes.NewReader(content))
   138  		if err != nil {
   139  			b.Fatal(err)
   140  		}
   141  	}
   142  }