Пример #1
0
func unpackIndex(ctx context.Context, pack *indexpack.Archive, dir string) error {
	fetcher := pack.Fetcher(ctx)
	return pack.ReadUnits(ctx, formatKey, func(u interface{}) error {
		unit, ok := u.(*apb.CompilationUnit)
		if !ok {
			return fmt.Errorf("%T is not a CompilationUnit", u)
		}
		idx, err := kindex.FromUnit(unit, fetcher)
		if err != nil {
			return fmt.Errorf("error creating kindex: %v", err)
		}
		path := kindexPath(dir, idx)
		if !*quiet {
			log.Println("Writing compilation unit to", path)
		}
		f, err := vfs.Create(ctx, path)
		if err != nil {
			return fmt.Errorf("error creating output file: %v", err)
		}
		if _, err := idx.WriteTo(f); err != nil {
			f.Close() // try to close file before returning
			return fmt.Errorf("error writing output file: %v", err)
		}
		return f.Close()
	})
}
Пример #2
0
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)
}
Пример #3
0
// Start begins profiling the program, writing data to the file given with the
// --cpu_profile flag.  If --cpu_profile was not given, nothing happens.  Start
// must not be called again until Stop is called.  The profile data in the
// --cpu_profile flag is overwritten on each call to Start.
func Start(ctx context.Context) error {
	mu.Lock()
	defer mu.Unlock()

	if *profCPU != "" {
		if file != nil {
			return errors.New("profiling already started")
		}

		f, err := vfs.Create(ctx, *profCPU)
		if err != nil {
			return fmt.Errorf("error creating profile file %q: %v", *profCPU)
		}
		file = f
		pprof.StartCPUProfile(f)
	}
	return nil
}
Пример #4
0
func main() {
	flag.Parse()

	if len(flag.Args()) > 2 || (gs != nil && len(flag.Args()) > 1) {
		fmt.Fprintf(os.Stderr, "ERROR: too many arguments %v\n", flag.Args())
		flag.Usage()
		os.Exit(1)
	}

	if gs != nil {
		defer gsutil.LogClose(context.Background(), gs)
	}

	var in io.ReadCloser = os.Stdin
	if gs == nil && len(flag.Args()) > 0 {
		file, err := vfs.Open(context.Background(), flag.Arg(0))
		if err != nil {
			log.Fatalf("Failed to open input file %q: %v", flag.Arg(0), err)
		}
		defer file.Close()
		in = file
	}

	outIdx := 1
	if gs != nil {
		outIdx = 0
	}

	var out io.WriteCloser = os.Stdout
	if len(flag.Args()) > outIdx {
		file, err := vfs.Create(context.Background(), flag.Arg(outIdx))
		if err != nil {
			log.Fatalf("Failed to create output file %q: %v", flag.Arg(outIdx), err)
		}
		defer file.Close()
		out = file
	}

	var (
		entries      <-chan *spb.Entry
		reverseEdges int
		triples      int
	)

	if gs == nil {
		entries = stream.ReadEntries(in)
	} else {
		ch := make(chan *spb.Entry)
		entries = ch
		go func() {
			defer close(ch)
			if err := gs.Scan(context.Background(), &spb.ScanRequest{}, func(e *spb.Entry) error {
				ch <- e
				return nil
			}); err != nil {
				log.Fatalf("Error scanning graphstore: %v", err)
			}
		}()
	}

	for entry := range entries {
		if schema.EdgeDirection(entry.EdgeKind) == schema.Reverse && !*keepReverseEdges {
			reverseEdges++
			continue
		}

		t, err := toTriple(entry)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Fprintln(out, t)
		triples++
	}

	if !*quiet {
		if !*keepReverseEdges {
			log.Printf("Skipped %d reverse edges", reverseEdges)
		}
		log.Printf("Wrote %d triples", triples)
	}
}
Пример #5
0
func main() {
	flag.Parse()
	if gs == nil {
		flagutil.UsageError("missing --graphstore")
	} else if *shardsToFiles != "" && *shards <= 0 {
		flagutil.UsageError("--sharded_file and --shards must be given together")
	} else if *shards > 0 && len(flag.Args()) > 0 {
		flagutil.UsageError("--shards and giving tickets for reads are mutually exclusive")
	}

	ctx := context.Background()

	wr := delimited.NewWriter(os.Stdout)
	var total int64
	if *shards <= 0 {
		entryFunc := func(entry *spb.Entry) error {
			if *count {
				total++
				return nil
			}
			return wr.PutProto(entry)
		}
		if len(flag.Args()) > 0 {
			if *targetTicket != "" || *factPrefix != "" {
				log.Fatal("--target and --fact_prefix are unsupported when given tickets")
			}
			if err := readEntries(ctx, gs, entryFunc, *edgeKind, flag.Args()); err != nil {
				log.Fatal(err)
			}
		} else {
			if err := scanEntries(ctx, gs, entryFunc, *edgeKind, *targetTicket, *factPrefix); err != nil {
				log.Fatal(err)
			}
		}
		if *count {
			fmt.Println(total)
		}
		return
	}

	sgs, ok := gs.(graphstore.Sharded)
	if !ok {
		log.Fatalf("Sharding unsupported for given GraphStore type: %T", gs)
	} else if *shardIndex >= *shards {
		log.Fatalf("Invalid shard index for %d shards: %d", *shards, *shardIndex)
	}

	if *count {
		cnt, err := sgs.Count(ctx, &spb.CountRequest{Index: *shardIndex, Shards: *shards})
		if err != nil {
			log.Fatalf("ERROR: %v", err)
		}
		fmt.Println(cnt)
		return
	} else if *shardsToFiles != "" {
		var wg sync.WaitGroup
		wg.Add(int(*shards))
		for i := int64(0); i < *shards; i++ {
			go func(i int64) {
				defer wg.Done()
				path := fmt.Sprintf("%s-%.5d-of-%.5d", *shardsToFiles, i, *shards)
				f, err := vfs.Create(ctx, path)
				if err != nil {
					log.Fatalf("Failed to create file %q: %v", path, err)
				}
				defer f.Close()
				wr := delimited.NewWriter(f)
				if err := sgs.Shard(ctx, &spb.ShardRequest{
					Index:  i,
					Shards: *shards,
				}, func(entry *spb.Entry) error {
					return wr.PutProto(entry)
				}); err != nil {
					log.Fatalf("GraphStore shard scan error: %v", err)
				}
			}(i)
		}
		wg.Wait()
		return
	}

	if err := sgs.Shard(ctx, &spb.ShardRequest{
		Index:  *shardIndex,
		Shards: *shards,
	}, func(entry *spb.Entry) error {
		return wr.PutProto(entry)
	}); err != nil {
		log.Fatalf("GraphStore shard scan error: %v", err)
	}
}