func newService(t *testing.T, entries []*spb.Entry) *GraphStoreService { gs := inmemory.Create() for req := range graphstore.BatchWrites(channelEntries(entries), 64) { if err := gs.Write(ctx, req); err != nil { t.Fatalf("Failed to write entries: %v", err) } } return NewGraphStoreService(gs) }
func main() { log.SetPrefix("write_entries: ") flag.Parse() if *numWorkers < 1 { flagutil.UsageErrorf("Invalid number of --workers %d (must be ≥ 1)", *numWorkers) } else if *batchSize < 1 { flagutil.UsageErrorf("Invalid --batch_size %d (must be ≥ 1)", *batchSize) } else if gs == nil { flagutil.UsageError("Missing --graphstore") } defer gsutil.LogClose(context.Background(), gs) gsutil.EnsureGracefulExit(gs) if *profCPU != "" { f, err := vfs.Create(context.Background(), *profCPU) if err != nil { log.Fatal(err) } defer f.Close() pprof.StartCPUProfile(f) defer pprof.StopCPUProfile() } writes := graphstore.BatchWrites(stream.ReadEntries(os.Stdin), *batchSize) var ( wg sync.WaitGroup numEntries uint64 ) wg.Add(*numWorkers) for i := 0; i < *numWorkers; i++ { go func() { defer wg.Done() num, err := writeEntries(context.Background(), gs, writes) if err != nil { log.Fatal(err) } atomic.AddUint64(&numEntries, num) }() } wg.Wait() log.Printf("Wrote %d entries", numEntries) }