github.com/tickoalcantara12/micro/v3@v3.0.0-20221007104245-9d75b9bcbab9/cmd/usage/plugin.go (about)

     1  package usage
     2  
     3  import (
     4  	"math/rand"
     5  	"net/http"
     6  	"os"
     7  	"sync/atomic"
     8  	"time"
     9  
    10  	"github.com/tickoalcantara12/micro/v3/plugin"
    11  	"github.com/tickoalcantara12/micro/v3/service/registry"
    12  	"github.com/tickoalcantara12/micro/v3/util/backoff"
    13  	"github.com/urfave/cli/v2"
    14  )
    15  
    16  func init() {
    17  	plugin.Register(Plugin())
    18  }
    19  
    20  func Plugin() plugin.Plugin {
    21  	var requests uint64
    22  
    23  	// create rand
    24  	source := rand.NewSource(time.Now().UnixNano())
    25  	r := rand.New(source)
    26  
    27  	return plugin.NewPlugin(
    28  		plugin.WithName("usage"),
    29  		plugin.WithInit(func(c *cli.Context) error {
    30  			// only do if enabled
    31  			if !c.Bool("report_usage") {
    32  				os.Setenv("MICRO_REPORT_USAGE", "false")
    33  				return nil
    34  			}
    35  
    36  			var service string
    37  
    38  			// set service name
    39  			if c.Args().Len() > 0 && len(c.Args().Get(0)) > 0 {
    40  				service = c.Args().Get(0)
    41  			}
    42  
    43  			// service subcommand
    44  			if service == "service" {
    45  				// set as the sub command
    46  				if v := c.Args().Get(1); len(v) > 0 {
    47  					service = v
    48  				}
    49  			}
    50  
    51  			// kick off the tracker
    52  			go func() {
    53  				// new report
    54  				u := New(service)
    55  
    56  				// initial publish in 30-60 seconds
    57  				d := 30 + r.Intn(30)
    58  				time.Sleep(time.Second * time.Duration(d))
    59  
    60  				for {
    61  					// get service list
    62  					s, _ := registry.ListServices()
    63  					// get requests
    64  					reqs := atomic.LoadUint64(&requests)
    65  					srvs := uint64(len(s))
    66  
    67  					// reset requests
    68  					atomic.StoreUint64(&requests, 0)
    69  
    70  					// set metrics
    71  					u.Metrics.Count["instances"] = uint64(1)
    72  					u.Metrics.Count["requests"] = reqs
    73  					u.Metrics.Count["services"] = srvs
    74  
    75  					// attempt to send report 3 times
    76  					for i := 1; i <= 3; i++ {
    77  						if err := Report(u); err != nil {
    78  							time.Sleep(backoff.Do(i * 2))
    79  							continue
    80  						}
    81  						break
    82  					}
    83  
    84  					// now sleep 24 hours
    85  					time.Sleep(time.Hour * 24)
    86  				}
    87  			}()
    88  
    89  			return nil
    90  		}),
    91  		plugin.WithHandler(func(h http.Handler) http.Handler {
    92  			// only enable if set
    93  			if v := os.Getenv("MICRO_REPORT_USAGE"); v == "false" {
    94  				return h
    95  			}
    96  
    97  			// return usage recorder
    98  			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    99  				// count requests
   100  				atomic.AddUint64(&requests, 1)
   101  				// serve the request
   102  				h.ServeHTTP(w, r)
   103  			})
   104  		}),
   105  	)
   106  }