github.com/distbuild/reclient@v0.0.0-20240401075343-3de72e395564/cmd/reproxystatus/main.go (about)

     1  // Copyright 2023 Google LLC
     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  // Binary to get status for running reproxy instances
    16  //
    17  // $ reproxystatus --color=
    18  package main
    19  
    20  import (
    21  	"context"
    22  	"flag"
    23  	"fmt"
    24  	"strings"
    25  	"time"
    26  
    27  	"github.com/fatih/color"
    28  	"github.com/gosuri/uilive"
    29  
    30  	"github.com/bazelbuild/reclient/internal/pkg/printer"
    31  	"github.com/bazelbuild/reclient/internal/pkg/rbeflag"
    32  	"github.com/bazelbuild/reclient/internal/pkg/reproxystatus"
    33  )
    34  
    35  var (
    36  	serverAddr = flag.String("server_address", "",
    37  		"The server address in the format of host:port for network, or unix:///file for unix domain sockets. "+
    38  			"If empty (default) then all reproxy instances will be dialed")
    39  	colorize = flag.String("color", "auto", "Control the output color mode; one of (off, on, auto)")
    40  	watch    = flag.Duration("watch", 0, "If greater than 0, every interval the ternimal will be cleared and the output will be printed.")
    41  )
    42  
    43  var (
    44  	dialTimeout = 30 * time.Second
    45  )
    46  
    47  func main() {
    48  	rbeflag.Parse()
    49  	switch strings.ToLower(*colorize) {
    50  	case "off", "false":
    51  		color.NoColor = true
    52  	case "on", "true":
    53  		color.NoColor = false
    54  	case "auto":
    55  	default:
    56  		fmt.Fprintf(color.Error, "ERROR: invalid --color mode: %v", *colorize)
    57  	}
    58  	if *watch < 0 {
    59  		printer.Fatal("ERROR: --watch needs to be >= 0")
    60  	}
    61  	ctx, cancel := context.WithCancel(context.Background())
    62  	defer cancel()
    63  	var tracker reproxystatus.ReproxyTracker
    64  	if *serverAddr != "" {
    65  		tracker = &reproxystatus.SingleReproxyTracker{ServerAddress: *serverAddr}
    66  	} else {
    67  		tracker = &reproxystatus.SocketReproxyTracker{}
    68  	}
    69  	if *watch == 0 {
    70  		tctx, cancel := context.WithTimeout(ctx, dialTimeout)
    71  		defer cancel()
    72  		reproxystatus.PrintSummaries(tctx, color.Output, tracker)
    73  	} else {
    74  		writer := uilive.New()
    75  		writer.Out = color.Output
    76  		ticker := time.NewTicker(*watch)
    77  		for ; true; <-ticker.C {
    78  			tctx, cancel := context.WithTimeout(ctx, dialTimeout)
    79  			reproxystatus.PrintSummaries(tctx, writer, tracker)
    80  			cancel()
    81  			writer.Flush()
    82  		}
    83  	}
    84  }