github.com/gobwas/gtrace@v0.4.3/examples/pinger/main.go (about)

     1  package main
     2  
     3  //go:generate gtrace
     4  
     5  import (
     6  	"bufio"
     7  	"context"
     8  	"flag"
     9  	"fmt"
    10  	"log"
    11  	"net/http"
    12  	"net/url"
    13  	"os"
    14  	"strings"
    15  	"sync"
    16  	"time"
    17  )
    18  
    19  type (
    20  	//gtrace:gen
    21  	//gtrace:set shortcut
    22  	//gtrace:set context
    23  	PingTrace struct {
    24  		OnRequest func(PingTraceRequestStart) func(PingTraceRequestDone)
    25  	}
    26  	PingTraceRequestStart struct {
    27  		Request *http.Request
    28  	}
    29  	PingTraceRequestDone struct {
    30  		Response *http.Response
    31  		Error    error
    32  	}
    33  )
    34  
    35  type Pinger struct {
    36  	Trace PingTrace
    37  
    38  	wg sync.WaitGroup
    39  }
    40  
    41  func (x *Pinger) Ping(ctx context.Context, s string) error {
    42  	u, err := url.ParseRequestURI(s)
    43  	if err != nil {
    44  		return err
    45  	}
    46  	req, err := http.NewRequest("HEAD", u.String(), nil)
    47  	if err != nil {
    48  		return err
    49  	}
    50  
    51  	x.wg.Add(1)
    52  	go func() {
    53  		defer x.wg.Done()
    54  
    55  		done := pingTraceOnRequest(ctx, x.Trace, req)
    56  		resp, err := http.DefaultClient.Do(req)
    57  		done(resp, err)
    58  	}()
    59  
    60  	return nil
    61  }
    62  
    63  func (x *Pinger) Close() {
    64  	x.wg.Wait()
    65  }
    66  
    67  func main() {
    68  	log.SetFlags(0)
    69  	log.SetPrefix("[logs] ")
    70  
    71  	var x Pinger
    72  	x.Trace = logTrace(x.Trace)
    73  	defer x.Close()
    74  
    75  	s := bufio.NewScanner(os.Stdin)
    76  
    77  	var (
    78  		prev     string
    79  		sampling int
    80  	)
    81  	flag.StringVar(&prev,
    82  		"u", "http://ya.ru",
    83  		"default url to ping",
    84  	)
    85  	flag.IntVar(&sampling,
    86  		"sampling", 5,
    87  		"sampling of stats probes",
    88  	)
    89  	flag.Parse()
    90  
    91  	fmt.Println("Welcome to pinger!")
    92  	fmt.Println("Default url is set to:", prev)
    93  	fmt.Println("Sampling for latency calculation is set to:", sampling)
    94  	fmt.Println("Type new url or press enter...")
    95  
    96  	for i := 0; s.Scan(); i++ {
    97  		ctx := context.Background()
    98  		if i%sampling == 0 {
    99  			ctx = WithPingTrace(ctx, statsTrace(PingTrace{}))
   100  		}
   101  		text := strings.TrimSpace(s.Text())
   102  		if text == "" {
   103  			text = prev
   104  		} else {
   105  			prev = text
   106  		}
   107  		err := x.Ping(ctx, text)
   108  		if err != nil {
   109  			log.Println(err)
   110  		}
   111  	}
   112  }
   113  
   114  func logTrace(s PingTrace) PingTrace {
   115  	return s.Compose(PingTrace{
   116  		OnRequest: func(req PingTraceRequestStart) func(PingTraceRequestDone) {
   117  			log.Printf("%s %s...", req.Request.Method, req.Request.URL)
   118  			return func(res PingTraceRequestDone) {
   119  				var (
   120  					status string
   121  					size   int64
   122  				)
   123  				if resp := res.Response; resp != nil {
   124  					status = resp.Status
   125  					size = resp.ContentLength
   126  				}
   127  				log.Printf(
   128  					"%s %s... done; status=%q size=%d err=%v",
   129  					req.Request.Method, req.Request.URL,
   130  					status, size, res.Error,
   131  				)
   132  			}
   133  		},
   134  	})
   135  }
   136  
   137  func statsTrace(s PingTrace) PingTrace {
   138  	return s.Compose(PingTrace{
   139  		OnRequest: func(req PingTraceRequestStart) func(PingTraceRequestDone) {
   140  			start := time.Now()
   141  			return func(res PingTraceRequestDone) {
   142  				log.Printf(
   143  					"%s %s    latency: %.2fms",
   144  					req.Request.Method, req.Request.URL,
   145  					time.Since(start).Seconds()*1000,
   146  				)
   147  			}
   148  		},
   149  	})
   150  }