// BatchWrites returns a channel of WriteRequests for the given entries.
// Consecutive entries with the same Source will be collected in the same
// WriteRequest, with each request containing up to maxSize updates.
func BatchWrites(entries <-chan *spb.Entry, maxSize int) <-chan *spb.WriteRequest {
ch := make(chan *spb.WriteRequest)
go func() {
defer close(ch)
var req *spb.WriteRequest
for entry := range entries {
update := &spb.WriteRequest_Update{
EdgeKind: entry.EdgeKind,
Target: entry.Target,
FactName: entry.FactName,
FactValue: entry.FactValue,
}
if req != nil && (!compare.VNamesEqual(req.Source, entry.Source) || len(req.Update) >= maxSize) {
ch <- req
req = nil
}
if req == nil {
req = &spb.WriteRequest{
Source: entry.Source,
Update: []*spb.WriteRequest_Update{update},
}
} else {
req.Update = append(req.Update, update)
}
}
if req != nil {
ch <- req
}
}()
return ch
}
func writeEdges(ctx context.Context, t table.Proto, edges <-chan *spb.Entry) error {
tempDir, err := ioutil.TempDir("", "reverse.edges")
if err != nil {
return fmt.Errorf("failed to create temporary directory: %v", err)
}
defer func() {
drainEntries(edges) // ensure channel is drained on errors
log.Println("Removing temporary edges table", tempDir)
if err := os.RemoveAll(tempDir); err != nil {
log.Printf("Failed to remove temporary directory %q: %v", tempDir, err)
}
}()
gs, err := leveldb.OpenGraphStore(tempDir, nil)
if err != nil {
return fmt.Errorf("failed to create temporary GraphStore: %v", err)
}
defer gs.Close(ctx)
log.Println("Writing temporary reverse edges table")
var writeReq *spb.WriteRequest
for e := range edges {
if writeReq != nil && !compare.VNamesEqual(e.Source, writeReq.Source) {
if err := writeWithReverses(ctx, gs, writeReq); err != nil {
return err
}
writeReq = nil
}
if writeReq == nil {
writeReq = &spb.WriteRequest{Source: e.Source}
}
writeReq.Update = append(writeReq.Update, &spb.WriteRequest_Update{
Target: e.Target,
EdgeKind: e.EdgeKind,
FactName: e.FactName,
FactValue: e.FactValue,
})
}
if writeReq != nil {
if err := writeWithReverses(ctx, gs, writeReq); err != nil {
return err
}
}
return writeEdgePages(ctx, t, gs)
}