// rebuildTable deletes the table if it exists, then creates the table, with the index column family.
func rebuildTable() error {
ctx, _ := context.WithTimeout(context.Background(), 5*time.Minute)
adminClient.DeleteTable(ctx, *tableName)
if err := adminClient.CreateTable(ctx, *tableName); err != nil {
return fmt.Errorf("CreateTable: %v", err)
}
time.Sleep(20 * time.Second)
if err := adminClient.CreateColumnFamily(ctx, *tableName, indexColumnFamily); err != nil {
return fmt.Errorf("CreateColumnFamily: %v", err)
}
if err := adminClient.CreateColumnFamily(ctx, *tableName, contentColumnFamily); err != nil {
return fmt.Errorf("CreateColumnFamily: %v", err)
}
// Open the prototype table. It contains a number of documents to get started with.
prototypeTable := client.Open(prototypeTableName)
var (
writeErr error // Set if any write fails.
mu sync.Mutex // Protects writeErr
wg sync.WaitGroup // Used to wait for all writes to finish.
)
copyRowToTable := func(row bigtable.Row) bool {
mu.Lock()
failed := writeErr != nil
mu.Unlock()
if failed {
return false
}
mut := bigtable.NewMutation()
for family, items := range row {
for _, item := range items {
// Get the column name, excluding the column family name and ':' character.
columnWithoutFamily := item.Column[len(family)+1:]
mut.Set(family, columnWithoutFamily, bigtable.Now(), item.Value)
}
}
wg.Add(1)
go func() {
// TODO: should use a semaphore to limit the number of concurrent writes.
if err := table.Apply(ctx, row.Key(), mut); err != nil {
mu.Lock()
writeErr = err
mu.Unlock()
}
wg.Done()
}()
return true
}
// Create a filter that only accepts the column families we're interested in.
filter := bigtable.FamilyFilter(indexColumnFamily + "|" + contentColumnFamily)
// Read every row from prototypeTable, and call copyRowToTable to copy it to our table.
err := prototypeTable.ReadRows(ctx, bigtable.InfiniteRange(""), copyRowToTable, bigtable.RowFilter(filter))
wg.Wait()
if err != nil {
return err
}
return writeErr
}
func doSet(ctx context.Context, args ...string) {
if len(args) < 3 {
log.Fatalf("usage: cbt set <table> <row> family:[column]=val[@ts] ...")
}
tbl := getClient().Open(args[0])
row := args[1]
mut := bigtable.NewMutation()
for _, arg := range args[2:] {
m := setArg.FindStringSubmatch(arg)
if m == nil {
log.Fatalf("Bad set arg %q", arg)
}
val := m[3]
ts := bigtable.Now()
if i := strings.LastIndex(val, "@"); i >= 0 {
// Try parsing a timestamp.
n, err := strconv.ParseInt(val[i+1:], 0, 64)
if err == nil {
val = val[:i]
ts = bigtable.Timestamp(n)
}
}
mut.Set(m[1], m[2], ts, []byte(val))
}
if err := tbl.Apply(ctx, row, mut); err != nil {
log.Fatalf("Applying mutation: %v", err)
}
}
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)
}
log.Printf("Dialing connections...")
client, err = bigtable.NewClient(context.Background(), config.Project, config.Zone, config.Cluster)
if err != nil {
log.Fatalf("Making bigtable.Client: %v", err)
}
defer client.Close()
adminClient, err = bigtable.NewAdminClient(context.Background(), config.Project, config.Zone, config.Cluster)
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, 100) // 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()
log.Printf("Reads (%d ok / %d tries):\n%v", reads.ok, reads.tries, newAggregate(reads.ds))
log.Printf("Writes (%d ok / %d tries):\n%v", writes.ok, writes.tries, newAggregate(writes.ds))
}
// handleAddDoc adds a document to the index.
func handleAddDoc(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
http.Error(w, "POST requests only", http.StatusMethodNotAllowed)
return
}
ctx, _ := context.WithTimeout(context.Background(), time.Minute)
name := r.FormValue("name")
if len(name) == 0 {
http.Error(w, "Empty document name!", http.StatusBadRequest)
return
}
content := r.FormValue("content")
if len(content) == 0 {
http.Error(w, "Empty document content!", http.StatusBadRequest)
return
}
var (
writeErr error // Set if any write fails.
mu sync.Mutex // Protects writeErr
wg sync.WaitGroup // Used to wait for all writes to finish.
)
// writeOneColumn writes one column in one row, updates err if there is an error,
// and signals wg that one operation has finished.
writeOneColumn := func(row, family, column, value string, ts bigtable.Timestamp) {
mut := bigtable.NewMutation()
mut.Set(family, column, ts, []byte(value))
err := table.Apply(ctx, row, mut)
if err != nil {
mu.Lock()
writeErr = err
mu.Unlock()
}
}
// Start a write to store the document content.
wg.Add(1)
go func() {
writeOneColumn(name, contentColumnFamily, "", content, bigtable.Now())
wg.Done()
}()
// Start writes to store the document name in the index for each word in the document.
words := tokenize(content)
for _, word := range words {
var (
row = word
family = indexColumnFamily
column = name
value = ""
ts = bigtable.Now()
)
wg.Add(1)
go func() {
// TODO: should use a semaphore to limit the number of concurrent writes.
writeOneColumn(row, family, column, value, ts)
wg.Done()
}()
}
wg.Wait()
if writeErr != nil {
http.Error(w, "Error writing to Bigtable: "+writeErr.Error(), http.StatusInternalServerError)
return
}
var buf bytes.Buffer
if err := addTemplate.ExecuteTemplate(&buf, "", struct{ Title string }{name}); err != nil {
http.Error(w, "Error executing HTML template: "+err.Error(), http.StatusInternalServerError)
return
}
io.Copy(w, &buf)
}