github.com/micro/go-micro/examples@v0.0.0-20210105173217-bf4ab679e18b/kubernetes/cmd/probe/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"os/signal"
     7  	"time"
     8  
     9  	"github.com/micro/cli/v2"
    10  	"github.com/micro/go-micro/v2/client"
    11  	gcli "github.com/micro/go-micro/v2/client/grpc"
    12  	"github.com/micro/go-micro/v2/config/cmd"
    13  	proto "github.com/micro/go-micro/v2/debug/service/proto"
    14  	"github.com/micro/go-micro/v2/util/log"
    15  	_ "github.com/micro/go-plugins/registry/kubernetes/v2"
    16  )
    17  
    18  const (
    19  	// StatusInvalidArguments indicates specified invalid arguments.
    20  	StatusInvalidArguments = 1
    21  	// StatusConnectionFailure indicates connection failed.
    22  	StatusConnectionFailure = 2
    23  	// StatusUnhealthy indicates rpc succeeded but indicates unhealthy service.
    24  	StatusUnhealthy = 4
    25  )
    26  
    27  var (
    28  	serverAddress string
    29  	serverName    string
    30  	connTimeout   time.Duration
    31  	rpcTimeout    time.Duration
    32  	verbose       bool
    33  )
    34  
    35  func init() {
    36  	os.Setenv("MICRO_REGISTRY", "kubernetes")
    37  	client.DefaultClient = gcli.NewClient()
    38  
    39  	argError := func(s string, v ...interface{}) {
    40  		log.Logf("error: "+s, v...)
    41  		os.Exit(StatusInvalidArguments)
    42  	}
    43  
    44  	app := cmd.App()
    45  	app.Flags = append(app.Flags,
    46  		&cli.DurationFlag{
    47  			Name:        "connect_timeout",
    48  			Value:       time.Second,
    49  			Usage:       "timeout for establishing connection",
    50  			EnvVars:     []string{"MICRO_CONNECT_TIMEOUT"},
    51  			Destination: &connTimeout,
    52  		},
    53  		&cli.DurationFlag{
    54  			Name:        "rpc_timeout",
    55  			Value:       time.Second,
    56  			Usage:       "timeout for health check rpc",
    57  			EnvVars:     []string{"MICRO_RPC_TIMEOUT"},
    58  			Destination: &rpcTimeout,
    59  		},
    60  		&cli.BoolFlag{
    61  			Name:        "verbose",
    62  			Usage:       "verbose logs",
    63  			EnvVars:     []string{"MICRO_VERBOSE"},
    64  			Destination: &verbose,
    65  		},
    66  	)
    67  
    68  	app.Action = func(c *cli.Context) error {
    69  		serverName = c.String("server_name")
    70  		serverAddress = c.String("server_address")
    71  
    72  		if len(serverName) == 0 {
    73  			argError("server name not set")
    74  		}
    75  		if len(serverAddress) == 0 {
    76  			argError("server address not set")
    77  		}
    78  		if connTimeout <= 0 {
    79  			argError("connection timeout must be greater than zero (specified: %v)", connTimeout)
    80  		}
    81  		if rpcTimeout <= 0 {
    82  			argError("rpc timeout must be greater than zero (specified: %v)", rpcTimeout)
    83  		}
    84  		return nil
    85  	}
    86  
    87  	cmd.Init()
    88  }
    89  
    90  func main() {
    91  	ctx, cancel := context.WithCancel(context.Background())
    92  
    93  	c := make(chan os.Signal, 1)
    94  	signal.Notify(c, os.Interrupt)
    95  	go func() {
    96  		sig := <-c
    97  		if sig == os.Interrupt {
    98  			log.Log("cancellation received")
    99  			cancel()
   100  			return
   101  		}
   102  	}()
   103  
   104  	if !verbose {
   105  		log.Log("establishing connection")
   106  	}
   107  
   108  	req := client.NewRequest(serverName, "Debug.Health", &proto.HealthRequest{})
   109  	rsp := &proto.HealthResponse{}
   110  	startTime := time.Now()
   111  
   112  	err := client.Call(ctx, req, rsp, client.WithAddress(serverAddress), client.WithDialTimeout(connTimeout), client.WithRequestTimeout(connTimeout))
   113  	if err != nil {
   114  		if err == context.DeadlineExceeded {
   115  			log.Logf("timeout: failed to connect service %q within %v", serverAddress, connTimeout)
   116  		} else {
   117  			log.Logf("error: failed to connect service at %q: %+v", serverAddress, err)
   118  		}
   119  		os.Exit(StatusConnectionFailure)
   120  	}
   121  
   122  	if !verbose {
   123  		log.Logf("time elapsed: %v", time.Since(startTime))
   124  	}
   125  
   126  	if rsp.Status != "ok" {
   127  		log.Logf("service unhealthy (responded with %q)", rsp.Status)
   128  		os.Exit(StatusUnhealthy)
   129  	}
   130  
   131  	log.Logf("status: %v", rsp.Status)
   132  }