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)
}
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
}
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
}
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
}
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
}
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
}