github.com/blend/go-sdk@v1.20220411.3/cmd/throttled_copy/main.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package main
     9  
    10  import (
    11  	"context"
    12  	"flag"
    13  	"fmt"
    14  	"io"
    15  	"os"
    16  	"time"
    17  
    18  	"github.com/blend/go-sdk/fileutil"
    19  	"github.com/blend/go-sdk/graceful"
    20  	"github.com/blend/go-sdk/ratelimiter"
    21  )
    22  
    23  var (
    24  	rateBytes   = flag.String("rate-bytes", "1024kb", "The throttle rate in bytes")
    25  	rateQuantum = flag.Duration("rate-quantum", time.Second, "The throttle quantum as a duration")
    26  	verbose     = flag.Bool("verbose", false, "If we should show verbose output")
    27  )
    28  
    29  func init() {
    30  	flag.Usage = func() {
    31  		fmt.Printf("throttled_copy SRC DST [flags]\n\nflags:\n")
    32  		flag.PrintDefaults()
    33  	}
    34  	flag.Parse()
    35  }
    36  
    37  func main() {
    38  	var src, dst string
    39  	if numArgs := len(flag.Args()); numArgs != 2 {
    40  		flag.Usage()
    41  		os.Exit(1)
    42  	}
    43  	src, dst = flag.Args()[0], flag.Args()[1]
    44  	maybeFatal(throttledCopy(graceful.Background(), src, dst))
    45  }
    46  
    47  func throttledCopy(ctx context.Context, src, dst string) error {
    48  	rateBytesValue, err := fileutil.ParseFileSize(*rateBytes)
    49  	if err != nil {
    50  		return err
    51  	}
    52  
    53  	var sr io.Reader
    54  	var dw io.Writer
    55  	if src == "-" {
    56  		sr = os.Stdin
    57  	} else {
    58  		sf, err := os.Open(src)
    59  		if err != nil {
    60  			return err
    61  		}
    62  		defer sf.Close()
    63  		sr = sf
    64  	}
    65  	if dst == "-" {
    66  		dw = os.Stdout
    67  	} else {
    68  		df, err := os.Create(dst)
    69  		if err != nil {
    70  			return err
    71  		}
    72  		defer df.Close()
    73  		dw = df
    74  	}
    75  
    76  	opts := []ratelimiter.CopyOption{
    77  		ratelimiter.OptCopyRateBytes(rateBytesValue),
    78  		ratelimiter.OptCopyRateQuantum(*rateQuantum),
    79  	}
    80  
    81  	if *verbose {
    82  		var written int64
    83  		opts = append(opts,
    84  			ratelimiter.OptCopyOnWrite(func(wr int, e time.Duration) {
    85  				written += int64(wr)
    86  				targetRate := float64(rateBytesValue) / (float64(*rateQuantum) / float64(time.Second))
    87  				rate := float64(wr) / (float64(e) / float64(time.Second))
    88  				fmt.Printf("written: %v, target: %0.2f/s, last: %0.2f/s\n", fileutil.FormatFileSize(written), targetRate, rate)
    89  			}),
    90  		)
    91  	}
    92  	_, err = ratelimiter.Copy(ctx, dw, sr, opts...)
    93  	return err
    94  }
    95  
    96  func maybeFatal(err error) {
    97  	if err != nil && err != context.Canceled {
    98  		fmt.Fprintf(os.Stderr, "%+v\n", err)
    99  		os.Exit(1)
   100  	}
   101  }