diff --git a/tools/batchgen/cmd/main.go b/tools/batchgen/cmd/main.go index 4428faa810..c90b588bd2 100644 --- a/tools/batchgen/cmd/main.go +++ b/tools/batchgen/cmd/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "io/ioutil" "log" "os" "strconv" @@ -28,32 +29,50 @@ func main() { log.Println("failed to convert factor to int:", err) return } - req, _, _, _, _ := makeStoreChunksRequest(100, 10, factor) - hosts := os.Args[2:] // assuming multiple hosts are passed as command line arguments + threads, err := strconv.Atoi(os.Args[2]) + if err != nil { + log.Println("failed to convert threads to int:", err) + return + } + hosts := os.Args[3:] // assuming multiple hosts are passed as command line arguments results := make(chan *pb.StoreChunksReply, len(hosts)) // channel to collect results errors := make(chan error, len(hosts)) // channel to collect errors + req, _, _, _, _ := makeStoreChunksRequest(100, 10, factor) + data, err := proto.Marshal(req) + if err != nil { + log.Fatalf("Failed to encode protobuf message: %v", err) + } + + // Write the data to a file + err = ioutil.WriteFile("req.pb", data, 0644) + if err != nil { + log.Fatalf("Failed to write to file: %v", err) + } + for _, host := range hosts { - go func(host string) { - conn, err := grpc.Dial(host, grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - errors <- err - return - } - defer conn.Close() - - client := pb.NewDispersalClient(conn) - opt := grpc.MaxCallSendMsgSize(60 * 1024 * 1024 * 1024) - - log.Println("sending chunks to operator", host, "request message size", proto.Size(req)) - reply, err := client.StoreChunks(context.Background(), req, opt) - if err != nil { - errors <- err - } else { - results <- reply - } - }(host) + for i := 0; i < threads; i++ { + go func(host string, i int) { + conn, err := grpc.Dial(host, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + errors <- err + return + } + defer conn.Close() + + client := pb.NewDispersalClient(conn) + opt := grpc.MaxCallSendMsgSize(60 * 1024 * 1024 * 1024) + + log.Println("host", host, "thread", i, "size", proto.Size(req)) + reply, err := client.StoreChunks(context.Background(), req, opt) + if err != nil { + errors <- err + } else { + results <- reply + } + }(host, i) + } } // Wait for all goroutines to finish and collect results