// NewClient creates a new Client for a given project and instance.
func NewClient(ctx context.Context, project, instance string, opts ...option.ClientOption) (*Client, error) {
o, err := btopt.DefaultClientOptions(prodAddr, Scope, clientUserAgent)
if err != nil {
return nil, err
}
// Default to a small connection pool that can be overridden.
o = append(o, option.WithGRPCConnectionPool(4))
o = append(o, opts...)
conn, err := transport.DialGRPC(ctx, o...)
if err != nil {
return nil, fmt.Errorf("dialing: %v", err)
}
return &Client{
conn: conn,
client: btpb.NewBigtableClient(conn),
project: project,
instance: instance,
}, nil
}
// WithGRPCConnectionPool returns a ClientOption that creates a pool of gRPC
// connections that requests will be balanced between.
// This is an EXPERIMENTAL API and may be changed or removed in the future.
func WithGRPCConnectionPool(size int) ClientOption {
return wrapOpt{option.WithGRPCConnectionPool(size)}
}
func main() {
var err error
config, err = cbtrc.Load()
if err != nil {
log.Fatal(err)
}
config.RegisterFlags()
flag.Parse()
if err := config.CheckFlags(); err != nil {
log.Fatal(err)
}
if config.Creds != "" {
os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", config.Creds)
}
if flag.NArg() != 0 {
flag.Usage()
os.Exit(1)
}
var options []option.ClientOption
if *poolSize > 1 {
options = append(options, option.WithGRPCConnectionPool(*poolSize))
}
var csvFile *os.File
if *csvOutput != "" {
csvFile, err = os.Create(*csvOutput)
if err != nil {
log.Fatalf("creating csv output file: %v", err)
}
defer csvFile.Close()
log.Printf("Writing statistics to %q ...", *csvOutput)
}
log.Printf("Dialing connections...")
client, err = bigtable.NewClient(context.Background(), config.Project, config.Instance, options...)
if err != nil {
log.Fatalf("Making bigtable.Client: %v", err)
}
defer client.Close()
adminClient, err = bigtable.NewAdminClient(context.Background(), config.Project, config.Instance)
if err != nil {
log.Fatalf("Making bigtable.AdminClient: %v", err)
}
defer adminClient.Close()
// Create a scratch table.
log.Printf("Setting up scratch table...")
if err := adminClient.CreateTable(context.Background(), *scratchTable); err != nil {
log.Fatalf("Making scratch table %q: %v", *scratchTable, err)
}
if err := adminClient.CreateColumnFamily(context.Background(), *scratchTable, "f"); err != nil {
log.Fatalf("Making scratch table column family: %v", err)
}
// Upon a successful run, delete the table. Don't bother checking for errors.
defer adminClient.DeleteTable(context.Background(), *scratchTable)
log.Printf("Starting load test... (run for %v)", *runFor)
tbl := client.Open(*scratchTable)
sem := make(chan int, *reqCount) // limit the number of requests happening at once
var reads, writes stats
stopTime := time.Now().Add(*runFor)
var wg sync.WaitGroup
for time.Now().Before(stopTime) {
sem <- 1
wg.Add(1)
go func() {
defer wg.Done()
defer func() { <-sem }()
ok := true
opStart := time.Now()
var stats *stats
defer func() {
stats.Record(ok, time.Since(opStart))
}()
row := fmt.Sprintf("row%d", rand.Intn(100)) // operate on 1 of 100 rows
switch rand.Intn(10) {
default:
// read
stats = &reads
_, err := tbl.ReadRow(context.Background(), row, bigtable.RowFilter(bigtable.LatestNFilter(1)))
if err != nil {
log.Printf("Error doing read: %v", err)
ok = false
}
case 0, 1, 2, 3, 4:
// write
stats = &writes
mut := bigtable.NewMutation()
mut.Set("f", "col", bigtable.Now(), bytes.Repeat([]byte("0"), 1<<10)) // 1 KB write
if err := tbl.Apply(context.Background(), row, mut); err != nil {
log.Printf("Error doing mutation: %v", err)
ok = false
}
}
}()
}
wg.Wait()
readsAgg := stat.NewAggregate("reads", reads.ds, reads.tries-reads.ok)
writesAgg := stat.NewAggregate("writes", writes.ds, writes.tries-writes.ok)
log.Printf("Reads (%d ok / %d tries):\n%v", reads.ok, reads.tries, readsAgg)
log.Printf("Writes (%d ok / %d tries):\n%v", writes.ok, writes.tries, writesAgg)
if csvFile != nil {
stat.WriteCSV([]*stat.Aggregate{readsAgg, writesAgg}, csvFile)
}
}