github.com/blend/go-sdk@v1.20220411.3/statsd/_bench/server/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  	"flag"
    12  	"fmt"
    13  	"io"
    14  	"log"
    15  	"os"
    16  	"sync"
    17  	"time"
    18  
    19  	"github.com/blend/go-sdk/statsd"
    20  )
    21  
    22  var (
    23  	bindAddr = flag.String("bind-addr", "127.0.0.1:0", "The bind address, defautls to a random local port")
    24  	verbose  = flag.Bool("verbose", false, "If we should print each metric the server receives")
    25  )
    26  
    27  func main() {
    28  	flag.Parse()
    29  	log.SetOutput(logger{os.Stdout})
    30  
    31  	listener, err := statsd.NewUDPListener(*bindAddr)
    32  	if err != nil {
    33  		log.Fatal(err)
    34  	}
    35  
    36  	srv := &statsd.Server{
    37  		Listener: listener,
    38  		Handler:  handleMetrics,
    39  	}
    40  
    41  	go printMetricCounts()
    42  
    43  	log.Printf("server listenening on: %s", listener.LocalAddr().String())
    44  	if err := srv.Start(); err != nil {
    45  		log.Fatal(err)
    46  	}
    47  }
    48  
    49  var (
    50  	metricRates   = map[string]*rate{}
    51  	metricRatesMu sync.Mutex
    52  )
    53  
    54  func printMetricCounts() {
    55  	for {
    56  		<-time.Tick(10 * time.Second)
    57  
    58  		log.Println("---")
    59  		log.Println("Metric Stats:")
    60  		for key, rate := range metricRates {
    61  			log.Printf("%s: %s\n", key, rate.String())
    62  		}
    63  	}
    64  }
    65  
    66  func handleMetrics(ms ...statsd.Metric) {
    67  	metricRatesMu.Lock()
    68  	defer metricRatesMu.Unlock()
    69  
    70  	for _, m := range ms {
    71  		if *verbose {
    72  			log.Printf("%#v\n", m)
    73  		}
    74  		_, ok := metricRates[m.Name]
    75  		if !ok {
    76  			metricRates[m.Name] = &rate{
    77  				Count: 1,
    78  				Time:  time.Now(),
    79  			}
    80  			return
    81  		}
    82  		metricRates[m.Name].Count++
    83  	}
    84  }
    85  
    86  // rate
    87  type rate struct {
    88  	Count int
    89  	Time  time.Time
    90  }
    91  
    92  func (r rate) String() string {
    93  	if r.Count == 0 {
    94  		return "N/A"
    95  	}
    96  	quantum := float64(time.Since(r.Time)) / float64(time.Second)
    97  	rate := float64(r.Count) / quantum
    98  	return fmt.Sprintf("%d %0.2f/s", r.Count, rate)
    99  }
   100  
   101  type logger struct {
   102  	wr io.Writer
   103  }
   104  
   105  func (l logger) Write(contents []byte) (int, error) {
   106  	return fmt.Fprint(l.wr, time.Now().UTC().Format(time.RFC3339Nano), " ", string(contents))
   107  }