Example #1
0
func writeFileTree(ctx context.Context, t table.Proto, files <-chan *spb.VName) error {
	tree := filetree.NewMap()
	for f := range files {
		tree.AddFile(f)
		// TODO(schroederc): evict finished directories (based on GraphStore order)
	}

	for corpus, roots := range tree.M {
		for root, dirs := range roots {
			for path, dir := range dirs {
				if err := t.Put(ftsrv.DirKey(corpus, root, path), dir); err != nil {
					return err
				}
			}
		}
	}
	cr, err := tree.CorpusRoots(ctx, &ftpb.CorpusRootsRequest{})
	if err != nil {
		return err
	}
	return t.Put(ftsrv.CorpusRootsKey, cr)
}
Example #2
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
}
Example #3
0
func main() {
	flag.Parse()
	if *servingTable == "" && gs == nil {
		flagutil.UsageError("missing either --serving_table or --graphstore")
	} else if *httpListeningAddr == "" && *grpcListeningAddr == "" {
		flagutil.UsageError("missing either --listen or --grpc_listen argument")
	} else if *servingTable != "" && gs != nil {
		flagutil.UsageError("--serving_table and --graphstore are mutually exclusive")
	}

	var (
		xs xrefs.Service
		ft filetree.Service
		sr search.Service
	)

	ctx := context.Background()
	if *servingTable != "" {
		db, err := leveldb.Open(*servingTable, nil)
		if err != nil {
			log.Fatalf("Error opening db at %q: %v", *servingTable, err)
		}
		defer db.Close()
		tbl := table.ProtoBatchParallel{&table.KVProto{db}}
		xs = xsrv.NewCombinedTable(tbl)
		ft = &ftsrv.Table{tbl}
		sr = &srchsrv.Table{&table.KVInverted{db}}
	} else {
		log.Println("WARNING: serving directly from a GraphStore can be slow; you may want to use a --serving_table")
		if f, ok := gs.(filetree.Service); ok {
			log.Printf("Using %T directly as filetree service", gs)
			ft = f
		} else {
			m := filetree.NewMap()
			if err := m.Populate(ctx, gs); err != nil {
				log.Fatalf("Error populating file tree from GraphStore: %v", err)
			}
			ft = m
		}

		if x, ok := gs.(xrefs.Service); ok {
			log.Printf("Using %T directly as xrefs service", gs)
			xs = x
		} else {
			if err := xstore.EnsureReverseEdges(ctx, gs); err != nil {
				log.Fatalf("Error ensuring reverse edges in GraphStore: %v", err)
			}
			xs = xstore.NewGraphStoreService(gs)
		}

		if s, ok := gs.(search.Service); ok {
			log.Printf("Using %T directly as search service", gs)
			sr = s
		}
	}

	if sr == nil {
		log.Println("Search API not supported")
	}

	if *grpcListeningAddr != "" {
		srv := grpc.NewServer()
		xpb.RegisterXRefServiceServer(srv, xs)
		ftpb.RegisterFileTreeServiceServer(srv, ft)
		if sr != nil {
			spb.RegisterSearchServiceServer(srv, sr)
		}
		go startGRPC(srv)
	}

	if *httpListeningAddr != "" {
		xrefs.RegisterHTTPHandlers(ctx, xs, http.DefaultServeMux)
		filetree.RegisterHTTPHandlers(ctx, ft, http.DefaultServeMux)
		if sr != nil {
			search.RegisterHTTPHandlers(ctx, sr, http.DefaultServeMux)
		}
		go startHTTP()
	}

	select {} // block forever
}
Example #4
0
func main() {
	flag.Parse()
	if *servingTable == "" && gs == nil {
		flagutil.UsageError("missing either --serving_table or --graphstore")
	} else if *httpListeningAddr == "" && *grpcListeningAddr == "" && *tlsListeningAddr == "" {
		flagutil.UsageError("missing either --listen, --tls_listen, or --grpc_listen argument")
	} else if *servingTable != "" && gs != nil {
		flagutil.UsageError("--serving_table and --graphstore are mutually exclusive")
	} else if *tlsListeningAddr != "" && (*tlsCertFile == "" || *tlsKeyFile == "") {
		flagutil.UsageError("--tls_cert_file and --tls_key_file are required if given --tls_listen")
	} else if flag.NArg() > 0 {
		flagutil.UsageErrorf("unknown non-flag arguments given: %v", flag.Args())
	}

	var (
		xs xrefs.Service
		ft filetree.Service
		sr search.Service
	)

	ctx := context.Background()
	if *servingTable != "" {
		db, err := leveldb.Open(*servingTable, nil)
		if err != nil {
			log.Fatalf("Error opening db at %q: %v", *servingTable, err)
		}
		defer db.Close()
		tbl := table.ProtoBatchParallel{&table.KVProto{db}}
		xs = xsrv.NewCombinedTable(tbl)
		ft = &ftsrv.Table{tbl}
		sr = &srchsrv.Table{&table.KVInverted{db}}
	} else {
		log.Println("WARNING: serving directly from a GraphStore can be slow; you may want to use a --serving_table")
		if f, ok := gs.(filetree.Service); ok {
			log.Printf("Using %T directly as filetree service", gs)
			ft = f
		} else {
			m := filetree.NewMap()
			if err := m.Populate(ctx, gs); err != nil {
				log.Fatalf("Error populating file tree from GraphStore: %v", err)
			}
			ft = m
		}

		if x, ok := gs.(xrefs.Service); ok {
			log.Printf("Using %T directly as xrefs service", gs)
			xs = x
		} else {
			if err := xstore.EnsureReverseEdges(ctx, gs); err != nil {
				log.Fatalf("Error ensuring reverse edges in GraphStore: %v", err)
			}
			xs = xstore.NewGraphStoreService(gs)
		}

		if s, ok := gs.(search.Service); ok {
			log.Printf("Using %T directly as search service", gs)
			sr = s
		}
	}

	if sr == nil {
		log.Println("Search API not supported")
	}

	if *grpcListeningAddr != "" {
		srv := grpc.NewServer()
		xpb.RegisterXRefServiceServer(srv, xs)
		ftpb.RegisterFileTreeServiceServer(srv, ft)
		if sr != nil {
			spb.RegisterSearchServiceServer(srv, sr)
		}
		go startGRPC(srv)
	}

	if *httpListeningAddr != "" || *tlsListeningAddr != "" {
		apiMux := http.NewServeMux()
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			if *httpAllowOrigin != "" {
				w.Header().Set("Access-Control-Allow-Origin", *httpAllowOrigin)
			}
			apiMux.ServeHTTP(w, r)
		})

		xrefs.RegisterHTTPHandlers(ctx, xs, apiMux)
		filetree.RegisterHTTPHandlers(ctx, ft, apiMux)
		if sr != nil {
			search.RegisterHTTPHandlers(ctx, sr, apiMux)
		}
		if *publicResources != "" {
			log.Println("Serving public resources at", *publicResources)
			if s, err := os.Stat(*publicResources); err != nil {
				log.Fatalf("ERROR: could not get FileInfo for %q: %v", *publicResources, err)
			} else if !s.IsDir() {
				log.Fatalf("ERROR: %q is not a directory", *publicResources)
			}
			apiMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
				http.ServeFile(w, r, filepath.Join(*publicResources, filepath.Clean(r.URL.Path)))
			})
		}
	}
	if *httpListeningAddr != "" {
		go startHTTP()
	}
	if *tlsListeningAddr != "" {
		go startTLS()
	}

	select {} // block forever
}
Example #5
0
func combineNodesAndEdges(ctx context.Context, opts *Options, out *servingOutput, rdIn stream.EntryReader) (disksort.Interface, error) {
	log.Println("Writing partial edges")

	tree := filetree.NewMap()
	rd := func(f func(*spb.Entry) error) error {
		return rdIn(func(e *spb.Entry) error {
			if e.FactName == facts.NodeKind && string(e.FactValue) == nodes.File {
				tree.AddFile(e.Source)
				// TODO(schroederc): evict finished directories (based on GraphStore order)
			}
			return f(e)
		})
	}

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

	bIdx := out.idx.Buffered()
	if err := assemble.Sources(rd, func(src *ipb.Source) error {
		return writePartialEdges(ctx, partialSorter, bIdx, src)
	}); 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
			if e.Target != nil {
				if opts.Verbose {
					log.Printf("WARNING: missing node facts for: %q", e.Source.Ticket)
				}
				// This is needed to satisfy later parts of the pipeline that look for targetless edges
				// to signify new nodes.
				if err := cSorter.Add(&srvpb.Edge{Source: &srvpb.Node{Ticket: e.Source.Ticket}}); err != nil {
					return fmt.Errorf("error writing complete edge: %v", err)
				}
			}
		}
		if e.Target == nil {
			// pass-through self-edges
			return cSorter.Add(e)
		}
		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
}
Example #6
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
}