github.com/google/cloudprober@v0.11.3/examples/extensions/myprober/myprobe/myprobe.go (about)

     1  package myprobe
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"strconv"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/google/cloudprober/logger"
    12  	"github.com/google/cloudprober/metrics"
    13  	"github.com/google/cloudprober/probes/options"
    14  	"github.com/google/cloudprober/targets/endpoint"
    15  	"github.com/hoisie/redis"
    16  )
    17  
    18  // Probe holds aggregate information about all probe runs, per-target.
    19  type Probe struct {
    20  	name    string
    21  	c       *ProbeConf
    22  	targets []endpoint.Endpoint
    23  	opts    *options.Options
    24  
    25  	res map[string]*metrics.EventMetrics // Results by target
    26  	l   *logger.Logger
    27  }
    28  
    29  // Init initializes the probe with the given params.
    30  func (p *Probe) Init(name string, opts *options.Options) error {
    31  	c, ok := opts.ProbeConf.(*ProbeConf)
    32  	if !ok {
    33  		return fmt.Errorf("not a my probe config")
    34  	}
    35  	p.c = c
    36  	p.name = name
    37  	p.opts = opts
    38  	p.l = opts.Logger
    39  
    40  	p.res = make(map[string]*metrics.EventMetrics)
    41  	return nil
    42  }
    43  
    44  // Start starts and runs the probe indefinitely.
    45  func (p *Probe) Start(ctx context.Context, dataChan chan *metrics.EventMetrics) {
    46  	probeTicker := time.NewTicker(p.opts.Interval)
    47  
    48  	for {
    49  		select {
    50  		case <-ctx.Done():
    51  			probeTicker.Stop()
    52  			return
    53  		case <-probeTicker.C:
    54  			// On probe tick, write data to the channel and run probe.
    55  			for _, em := range p.res {
    56  				dataChan <- em.Clone()
    57  			}
    58  			p.targets = p.opts.Targets.ListEndpoints()
    59  			p.initProbeMetrics()
    60  			probeCtx, cancelFunc := context.WithDeadline(ctx, time.Now().Add(p.opts.Timeout))
    61  			p.runProbe(probeCtx)
    62  			cancelFunc()
    63  		}
    64  	}
    65  }
    66  
    67  // initProbeMetrics initializes missing probe metrics.
    68  func (p *Probe) initProbeMetrics() {
    69  	for _, target := range p.targets {
    70  		if p.res[target.Name] != nil {
    71  			continue
    72  		}
    73  		var latVal metrics.Value
    74  		if p.opts.LatencyDist != nil {
    75  			latVal = p.opts.LatencyDist.Clone()
    76  		} else {
    77  			latVal = metrics.NewFloat(0)
    78  		}
    79  		p.res[target.Name] = metrics.NewEventMetrics(time.Now()).
    80  			AddMetric("total", metrics.NewInt(0)).
    81  			AddMetric("success", metrics.NewInt(0)).
    82  			AddMetric("latency", latVal).
    83  			AddLabel("ptype", "redis").
    84  			AddLabel("probe", p.name).
    85  			AddLabel("dst", target.Name)
    86  	}
    87  }
    88  
    89  // runProbeForTarget runs probe for a single target.
    90  func (p *Probe) runProbeForTarget(ctx context.Context, target endpoint.Endpoint) error {
    91  	client := &redis.Client{
    92  		Addr: net.JoinHostPort(target.Name, strconv.Itoa(target.Port)),
    93  	}
    94  	key := p.c.GetKey()
    95  	val := p.c.GetValue()
    96  
    97  	switch p.c.GetOp() {
    98  	case ProbeConf_SET:
    99  		return client.Set(key, []byte(val))
   100  	case ProbeConf_GET:
   101  		_, err := client.Get(key)
   102  		return err
   103  	case ProbeConf_DELETE:
   104  		_, err := client.Del(key)
   105  		return err
   106  	default:
   107  		return fmt.Errorf("unknown op: %s", p.c.GetOp())
   108  	}
   109  }
   110  
   111  // runProbe runs probe for all targets and update EventMetrics.
   112  func (p *Probe) runProbe(ctx context.Context) {
   113  	p.targets = p.opts.Targets.ListEndpoints()
   114  
   115  	var wg sync.WaitGroup
   116  	for _, target := range p.targets {
   117  		wg.Add(1)
   118  
   119  		go func(target endpoint.Endpoint, em *metrics.EventMetrics) {
   120  			defer wg.Done()
   121  			start := time.Now()
   122  			em.Timestamp = start
   123  			em.Metric("total").AddInt64(1)
   124  			err := p.runProbeForTarget(ctx, target) // run probe just for a single target
   125  			if err != nil {
   126  				p.l.Errorf(err.Error())
   127  				return
   128  			}
   129  			em.Metric("success").AddInt64(1)
   130  			em.Metric("latency").AddFloat64(time.Now().Sub(start).Seconds() / p.opts.LatencyUnit.Seconds())
   131  		}(target, p.res[target.Name])
   132  	}
   133  
   134  	wg.Wait()
   135  }