gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/tools/tracereplay/main/main.go (about)

     1  // Copyright 2022 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package main implements a tool that can save and replay messages from
    16  // issued from remote.Remote.
    17  package main
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"os"
    23  	"os/signal"
    24  
    25  	"github.com/google/subcommands"
    26  	"gvisor.dev/gvisor/runsc/flag"
    27  	"gvisor.dev/gvisor/tools/tracereplay"
    28  )
    29  
    30  func main() {
    31  	subcommands.Register(subcommands.HelpCommand(), "")
    32  	subcommands.Register(subcommands.FlagsCommand(), "")
    33  	subcommands.Register(&saveCmd{}, "")
    34  	subcommands.Register(&replayCmd{}, "")
    35  	flag.CommandLine.Parse(os.Args[1:])
    36  	os.Exit(int(subcommands.Execute(context.Background())))
    37  }
    38  
    39  // saveCmd implements subcommands.Command for the "save" command.
    40  type saveCmd struct {
    41  	endpoint string
    42  	out      string
    43  	prefix   string
    44  }
    45  
    46  // Name implements subcommands.Command.
    47  func (*saveCmd) Name() string {
    48  	return "save"
    49  }
    50  
    51  // Synopsis implements subcommands.Command.
    52  func (*saveCmd) Synopsis() string {
    53  	return "save trace sessions to files"
    54  }
    55  
    56  // Usage implements subcommands.Command.
    57  func (*saveCmd) Usage() string {
    58  	return `save [flags] - save trace sessions to files
    59  `
    60  }
    61  
    62  // SetFlags implements subcommands.Command.
    63  func (c *saveCmd) SetFlags(f *flag.FlagSet) {
    64  	f.StringVar(&c.endpoint, "endpoint", "", "path to trace server endpoint to connect")
    65  	f.StringVar(&c.out, "out", "./replay", "path to a directory where trace files will be saved")
    66  	f.StringVar(&c.prefix, "prefix", "client-", "name to be prefixed to each trace file")
    67  }
    68  
    69  // Execute implements subcommands.Command.
    70  func (c *saveCmd) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcommands.ExitStatus {
    71  	if f.NArg() > 0 {
    72  		fmt.Fprintf(os.Stderr, "unexpected argument: %s\n", f.Args())
    73  		return subcommands.ExitUsageError
    74  	}
    75  	if len(c.endpoint) == 0 {
    76  		fmt.Fprintf(os.Stderr, "--endpoint is required\n")
    77  		return subcommands.ExitUsageError
    78  	}
    79  	_ = os.Remove(c.endpoint)
    80  
    81  	server := tracereplay.NewSave(c.endpoint, c.out, c.prefix)
    82  	defer server.Close()
    83  
    84  	if err := server.Start(); err != nil {
    85  		fmt.Fprintf(os.Stderr, "starting server: %v\n", err)
    86  		return subcommands.ExitFailure
    87  	}
    88  
    89  	ch := make(chan os.Signal, 1)
    90  	signal.Notify(ch, os.Interrupt)
    91  
    92  	done := make(chan struct{})
    93  	go func() {
    94  		<-ch
    95  		fmt.Printf("Ctrl-C pressed, stopping.\n")
    96  		done <- struct{}{}
    97  	}()
    98  
    99  	fmt.Printf("Listening on %q. Press ctrl-C to stop...\n", c.endpoint)
   100  	<-done
   101  	return subcommands.ExitSuccess
   102  }
   103  
   104  // replayCmd implements subcommands.Command for the "replay" command.
   105  type replayCmd struct {
   106  	endpoint string
   107  	in       string
   108  }
   109  
   110  // Name implements subcommands.Command.
   111  func (*replayCmd) Name() string {
   112  	return "replay"
   113  }
   114  
   115  // Synopsis implements subcommands.Command.
   116  func (*replayCmd) Synopsis() string {
   117  	return "replay a trace session from a file"
   118  }
   119  
   120  // Usage implements subcommands.Command.
   121  func (*replayCmd) Usage() string {
   122  	return `replay [flags] - replay a trace session from a file
   123  `
   124  }
   125  
   126  // SetFlags implements subcommands.Command.
   127  func (c *replayCmd) SetFlags(f *flag.FlagSet) {
   128  	f.StringVar(&c.endpoint, "endpoint", "", "path to trace server endpoint to connect")
   129  	f.StringVar(&c.in, "in", "", "path to trace file containing messages to be replayed")
   130  }
   131  
   132  // Execute implements subcommands.Command.
   133  func (c *replayCmd) Execute(_ context.Context, f *flag.FlagSet, _ ...any) subcommands.ExitStatus {
   134  	if f.NArg() > 0 {
   135  		fmt.Fprintf(os.Stderr, "unexpected argument: %s\n", f.Args())
   136  		return subcommands.ExitUsageError
   137  	}
   138  	if len(c.in) == 0 {
   139  		fmt.Fprintf(os.Stderr, "--in is required\n")
   140  		return subcommands.ExitUsageError
   141  	}
   142  
   143  	r := tracereplay.Replay{
   144  		Endpoint: c.endpoint,
   145  		In:       c.in,
   146  	}
   147  	if err := r.Execute(); err != nil {
   148  		fmt.Fprintln(os.Stderr, err)
   149  		return subcommands.ExitFailure
   150  	}
   151  	return subcommands.ExitSuccess
   152  }