agones.dev/agones@v1.54.0/cmd/ping/main.go (about)

     1  // Copyright 2018 Google LLC All Rights Reserved.
     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 for the pinger service for RTT measurement.
    16  package main
    17  
    18  import (
    19  	"context"
    20  	"net/http"
    21  	"strings"
    22  	"time"
    23  
    24  	"agones.dev/agones/pkg"
    25  	"agones.dev/agones/pkg/util/runtime"
    26  	"agones.dev/agones/pkg/util/signals"
    27  	"github.com/heptiolabs/healthcheck"
    28  	"github.com/pkg/errors"
    29  	"github.com/spf13/pflag"
    30  	"github.com/spf13/viper"
    31  	"golang.org/x/time/rate"
    32  )
    33  
    34  const (
    35  	httpResponseFlag = "http-response"
    36  	udpRateLimitFlag = "udp-rate-limit"
    37  )
    38  
    39  var (
    40  	logger = runtime.NewLoggerWithSource("main")
    41  )
    42  
    43  func main() {
    44  	ctlConf := parseEnvFlags()
    45  	if err := ctlConf.validate(); err != nil {
    46  		logger.WithError(err).Fatal("Could not create controller from environment or flags")
    47  	}
    48  
    49  	logger.WithField("version", pkg.Version).WithField("featureGates", runtime.EncodeFeatures()).
    50  		WithField("ctlConf", ctlConf).Info("Starting ping...")
    51  
    52  	ctx, _ := signals.NewSigKillContext()
    53  
    54  	udpSrv := serveUDP(ctx, ctlConf)
    55  	defer udpSrv.close()
    56  
    57  	h := healthcheck.NewHandler()
    58  	h.AddLivenessCheck("udp-server", udpSrv.Health)
    59  
    60  	cancel := serveHTTP(ctlConf, h)
    61  	defer cancel()
    62  
    63  	<-ctx.Done()
    64  	logger.Info("Shutting down...")
    65  }
    66  
    67  func serveUDP(ctx context.Context, ctlConf config) *udpServer {
    68  	s := newUDPServer(ctlConf.UDPRateLimit)
    69  	s.run(ctx)
    70  	return s
    71  }
    72  
    73  // serveHTTP starts the HTTP handler, and returns a cancel/shutdown function
    74  func serveHTTP(ctlConf config, h healthcheck.Handler) func() {
    75  	// we don't need a health checker, we already have a http endpoint that returns 200
    76  	mux := http.NewServeMux()
    77  	srv := &http.Server{
    78  		Addr:    ":8080",
    79  		Handler: mux,
    80  	}
    81  
    82  	// add health check as well
    83  	mux.HandleFunc("/live", h.LiveEndpoint)
    84  
    85  	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
    86  		if _, err := w.Write([]byte(ctlConf.HTTPResponse)); err != nil {
    87  			w.WriteHeader(http.StatusInternalServerError)
    88  			logger.WithError(err).Error("Error responding to http request")
    89  		}
    90  	})
    91  
    92  	go func() {
    93  		logger.Info("Starting HTTP Server...")
    94  		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
    95  			logger.WithError(err).Fatal("Could not start HTTP server")
    96  		}
    97  	}()
    98  
    99  	return func() {
   100  		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   101  		defer cancel()
   102  
   103  		if err := srv.Shutdown(ctx); err != nil {
   104  			logger.WithError(err).Fatal("Could not shut down HTTP server")
   105  		}
   106  		logger.Info("HTTP server was gracefully shut down")
   107  	}
   108  }
   109  
   110  // config retains the configuration information
   111  type config struct {
   112  	HTTPResponse string
   113  	UDPRateLimit rate.Limit
   114  }
   115  
   116  // validate returns an error if there is a validation problem
   117  func (c *config) validate() error {
   118  	if c.UDPRateLimit < 0 {
   119  		return errors.New("UDP Rate limit must be greater that or equal to zero")
   120  	}
   121  
   122  	return nil
   123  }
   124  
   125  func parseEnvFlags() config {
   126  	viper.SetDefault(httpResponseFlag, "ok")
   127  	viper.SetDefault(udpRateLimitFlag, 20)
   128  
   129  	pflag.String(httpResponseFlag, viper.GetString(httpResponseFlag), "Flag to set text value when a 200 response is returned. Can be useful to identify clusters. Defaults to 'ok' Can also use HTTP_RESPONSE env variable")
   130  	pflag.Float64(udpRateLimitFlag, viper.GetFloat64(httpResponseFlag), "Flag to set how many UDP requests can be handled by a single source IP per second. Defaults to 20. Can also use UDP_RATE_LIMIT env variable")
   131  	runtime.FeaturesBindFlags()
   132  	pflag.Parse()
   133  
   134  	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
   135  	runtime.Must(viper.BindEnv(httpResponseFlag))
   136  	runtime.Must(viper.BindEnv(udpRateLimitFlag))
   137  	runtime.Must(runtime.FeaturesBindEnv())
   138  
   139  	runtime.Must(runtime.ParseFeaturesFromEnv())
   140  
   141  	return config{
   142  		HTTPResponse: viper.GetString(httpResponseFlag),
   143  		UDPRateLimit: rate.Limit(viper.GetFloat64(udpRateLimitFlag)),
   144  	}
   145  }