kythe.io@v0.0.68-0.20240422202219-7225dbc01741/kythe/go/storage/tools/write_entries/write_entries.go (about)

     1  /*
     2   * Copyright 2014 The Kythe Authors. All rights reserved.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *   http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  // Binary write_entries reads a stream of protobuf entries on os.Stdin and
    18  // writes each to a graphstore server.
    19  //
    20  // Usage:
    21  //
    22  //	entry_emitter ... | write_entries --graphstore addr
    23  //
    24  // Example:
    25  //
    26  //	java_indexer_server --port 8181 &
    27  //	graphstore --port 9999 &
    28  //	analysis_driver --analyzer localhost:8181 /tmp/compilation.kzip | \
    29  //	  write_entries --workers 10 --graphstore localhost:9999
    30  //
    31  // Example:
    32  //
    33  //	zcat entries.gz | write_entries --graphstore gs/leveldb
    34  package main
    35  
    36  import (
    37  	"context"
    38  	"flag"
    39  	"os"
    40  	"sync"
    41  	"sync/atomic"
    42  
    43  	"kythe.io/kythe/go/services/graphstore"
    44  	"kythe.io/kythe/go/storage/gsutil"
    45  	"kythe.io/kythe/go/storage/stream"
    46  	"kythe.io/kythe/go/util/flagutil"
    47  	"kythe.io/kythe/go/util/log"
    48  	"kythe.io/kythe/go/util/profile"
    49  
    50  	spb "kythe.io/kythe/proto/storage_go_proto"
    51  
    52  	_ "kythe.io/kythe/go/services/graphstore/proxy"
    53  	_ "kythe.io/kythe/go/storage/leveldb"
    54  )
    55  
    56  var (
    57  	batchSize  = flag.Int("batch_size", 1024, "Maximum entries per write for consecutive entries with the same source")
    58  	numWorkers = flag.Int("workers", 1, "Number of concurrent workers writing to the GraphStore")
    59  
    60  	gs graphstore.Service
    61  )
    62  
    63  func init() {
    64  	flag.Usage = flagutil.SimpleUsage("Write a delimited stream of entries from stdin to a GraphStore",
    65  		"[--batch_size entries] [--workers n] --graphstore spec")
    66  	gsutil.Flag(&gs, "graphstore", "GraphStore to which to write the entry stream")
    67  }
    68  
    69  func main() {
    70  	flag.Parse()
    71  	if *numWorkers < 1 {
    72  		flagutil.UsageErrorf("Invalid number of --workers %d (must be ≥ 1)", *numWorkers)
    73  	} else if *batchSize < 1 {
    74  		flagutil.UsageErrorf("Invalid --batch_size %d (must be ≥ 1)", *batchSize)
    75  	} else if gs == nil {
    76  		flagutil.UsageError("Missing --graphstore")
    77  	}
    78  
    79  	ctx := context.Background()
    80  
    81  	defer gsutil.LogClose(ctx, gs)
    82  	gsutil.EnsureGracefulExit(gs)
    83  
    84  	if err := profile.Start(ctx); err != nil {
    85  		log.Fatal(err)
    86  	}
    87  	defer profile.Stop()
    88  
    89  	writes := graphstore.BatchWrites(stream.ReadEntries(os.Stdin), *batchSize)
    90  
    91  	var (
    92  		wg         sync.WaitGroup
    93  		numEntries uint64
    94  	)
    95  	wg.Add(*numWorkers)
    96  	for i := 0; i < *numWorkers; i++ {
    97  		go func() {
    98  			defer wg.Done()
    99  			num, err := writeEntries(ctx, gs, writes)
   100  			if err != nil {
   101  				log.Fatal(err)
   102  			}
   103  
   104  			atomic.AddUint64(&numEntries, num)
   105  		}()
   106  	}
   107  	wg.Wait()
   108  
   109  	log.InfoContextf(ctx, "Wrote %d entries", numEntries)
   110  }
   111  
   112  func writeEntries(ctx context.Context, s graphstore.Service, reqs <-chan *spb.WriteRequest) (uint64, error) {
   113  	var num uint64
   114  
   115  	for req := range reqs {
   116  		num += uint64(len(req.Update))
   117  		if err := s.Write(ctx, req); err != nil {
   118  			return 0, err
   119  		}
   120  	}
   121  
   122  	return num, nil
   123  }