overflow_probes/probe/probe.go

package probe

import (
	"context"
	"fmt"
	"net/http"

	"git.loafle.net/commons_go/logging"
	"git.loafle.net/overflow/overflow_probes/central/api/module"
	"git.loafle.net/overflow/overflow_probes/central/client"
	"git.loafle.net/overflow/overflow_probes/commons"
	"git.loafle.net/overflow/overflow_probes/config"
	opuu "git.loafle.net/overflow/overflow_probes/util/url"
)

const (
	probeEntryPoint  = "/probe"
	fileEntryPoint   = "/file"
	metricEntryPoint = "/metric"
)

type Prober interface {
	commons.Starter
	commons.Shutdowner
}

type probe struct {
	probeEntryURL  string
	fileEntryURL   string
	metricEntryURL string

	probeClient  client.Client
	fileClient   client.Client
	metricClient client.Client

	shutdown chan bool
}

func New() (Prober, error) {
	p := &probe{
		shutdown: make(chan bool),
	}

	var err error

	if p.probeEntryURL, err = opuu.Join(config.CFG.Central.URL, probeEntryPoint); nil != err {
		return nil, err
	}
	if p.fileEntryURL, err = opuu.Join(config.CFG.Central.URL, fileEntryPoint); nil != err {
		return nil, err
	}
	if p.metricEntryURL, err = opuu.Join(config.CFG.Central.URL, metricEntryPoint); nil != err {
		return nil, err
	}

	p.probeClient = client.New()
	p.fileClient = client.New()
	p.metricClient = client.New()

	return p, nil
}

func (p *probe) Start() error {

	return p.start()
}

func (p *probe) start() error {
	if err := p.connectToCentral(); nil != err {
		return err
	}

	p.listen()

	return nil
}

func (p *probe) listen() {
	go func() {
		for {
			select {
			case <-p.shutdown:
				break
			}
		}
	}()
}

func (p *probe) connectToCentral() error {
	var err error
	var res *http.Response
	if res, err = client.ConnectToCentralAsProbe(p.probeClient, p.probeEntryURL); nil != err {
		return err
	}

	encryptionKey := res.Header.Get(module.ProbeHeader_Probe_EncryptionKey)
	config.EncryptionKey = &encryptionKey

	p.probeClient.OnNotify(p.onNotify)

	if _, err = client.ConnectToCentralAsProbe(p.metricClient, p.metricEntryURL); nil != err {
		return err
	}

	return nil
}

func (p *probe) sendNotifyToCentral(method string, params ...string) {
	if err := p.probeClient.Notify(method, params); nil != err {
		logging.Logger.Error(fmt.Sprintf("Probe notify: %v", err))
	}
}

func (p *probe) Shutdown(ctx context.Context) error {
	for {
		p.stop(fmt.Errorf("Shutdown"))
		select {
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (p *probe) stop(err error) {
	defer close(p.shutdown)
}