Ejemplo n.º 1
0
func combineNodesAndEdges(ctx context.Context, opts *Options, out *servingOutput, gsEntries <-chan *spb.Entry) (disksort.Interface, error) {
	log.Println("Writing partial edges")

	tree := filetree.NewMap()

	partialSorter, err := opts.diskSorter(edgeLesser{}, edgeMarshaler{})
	if err != nil {
		return nil, err
	}

	bIdx := out.idx.Buffered()
	var src *spb.VName
	var entries []*spb.Entry
	for e := range gsEntries {
		if e.FactName == schema.NodeKindFact && string(e.FactValue) == schema.FileKind {
			tree.AddFile(e.Source)
			// TODO(schroederc): evict finished directories (based on GraphStore order)
		}

		if src == nil {
			src = e.Source
		} else if !compare.VNamesEqual(e.Source, src) {
			if err := writePartialEdges(ctx, partialSorter, bIdx, assemble.SourceFromEntries(entries)); err != nil {
				drainEntries(gsEntries)
				return nil, err
			}
			src = e.Source
			entries = nil
		}

		entries = append(entries, e)
	}
	if len(entries) > 0 {
		if err := writePartialEdges(ctx, partialSorter, bIdx, assemble.SourceFromEntries(entries)); err != nil {
			return nil, err
		}
	}
	if err := bIdx.Flush(ctx); err != nil {
		return nil, err
	}

	if err := writeFileTree(ctx, tree, out.xs); err != nil {
		return nil, fmt.Errorf("error writing file tree: %v", err)
	}
	tree = nil

	log.Println("Writing complete edges")

	cSorter, err := opts.diskSorter(edgeLesser{}, edgeMarshaler{})
	if err != nil {
		return nil, err
	}

	var n *srvpb.Node
	if err := partialSorter.Read(func(i interface{}) error {
		e := i.(*srvpb.Edge)
		if n == nil || n.Ticket != e.Source.Ticket {
			n = e.Source
			return cSorter.Add(e)
		} else if e.Target != nil {
			e.Source = n
			if err := writeCompletedEdges(ctx, cSorter, e); err != nil {
				return fmt.Errorf("error writing complete edge: %v", err)
			}
		}
		return nil
	}); err != nil {
		return nil, fmt.Errorf("error reading/writing edges: %v", err)
	}

	return cSorter, nil
}
Ejemplo n.º 2
0
func combineNodesAndEdges(ctx context.Context, out *servingOutput, gsEntries <-chan *spb.Entry) error {
	log.Println("Writing partial edges")

	tree := filetree.NewMap()

	var src *spb.VName
	var entries []*spb.Entry
	for e := range gsEntries {
		if e.FactName == schema.NodeKindFact && string(e.FactValue) == schema.FileKind {
			tree.AddFile(e.Source)
			// TODO(schroederc): evict finished directories (based on GraphStore order)
		}

		if src == nil {
			src = e.Source
		} else if !compare.VNamesEqual(e.Source, src) {
			if err := writePartialEdges(ctx, out, assemble.SourceFromEntries(entries)); err != nil {
				drainEntries(gsEntries)
				return err
			}
			src = e.Source
			entries = nil
		}

		entries = append(entries, e)
	}
	if len(entries) > 0 {
		if err := writePartialEdges(ctx, out, assemble.SourceFromEntries(entries)); err != nil {
			return err
		}
	}

	if err := writeFileTree(ctx, tree, out.xs); err != nil {
		return fmt.Errorf("error writing file tree: %v", err)
	}
	tree = nil

	log.Println("Writing complete edges")
	snapshot := out.completeEdges.NewSnapshot()
	defer snapshot.Close()
	it, err := out.completeEdges.ScanPrefix(nil, &keyvalue.Options{
		LargeRead: true,
		Snapshot:  snapshot,
	})
	if err != nil {
		return err
	}
	defer it.Close()

	var n *srvpb.Node
	var e srvpb.Edge
	for {
		k, v, err := it.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			return fmt.Errorf("error scanning partial edges table: %v", err)
		}

		ss := strings.Split(string(k), tempTableKeySep)
		if len(ss) != 3 {
			return fmt.Errorf("invalid partial edge table key: %q", string(k))
		}

		if err := proto.Unmarshal(v, &e); err != nil {
			return fmt.Errorf("invalid partial edge table value: %v", err)
		}

		if n == nil || n.Ticket != ss[0] {
			n = e.Source
		} else if e.Target != nil {
			e.Source = n
			if err := writeCompletedEdges(ctx, out.completeEdges, &e); err != nil {
				return fmt.Errorf("error writing complete edge: %v", err)
			}
		}
	}

	return nil
}