probe_gateway_metric/servlet/data-servlet.go

package servlet

import (
	"context"
	"fmt"

	"github.com/valyala/fasthttp"

	logging "git.loafle.net/commons/logging-go"
	"git.loafle.net/commons/server-go"
	"git.loafle.net/commons/server-go/socket"
	cssw "git.loafle.net/commons/server-go/socket/web"
	occp "git.loafle.net/overflow/commons-go/config/probe"
	og "git.loafle.net/overflow/gateway"
	"git.loafle.net/overflow/gateway/external/grpc"
	"git.loafle.net/overflow/gateway/external/kafka"
)

type DataServlet interface {
	cssw.Servlet
}

type DataServlets struct {
	cssw.Servlets
}

func (s *DataServlets) Init(serverCtx server.ServerCtx) error {
	if err := s.Servlets.Init(serverCtx); nil != err {
		return err
	}

	return nil
}

func (s *DataServlets) OnStart(serverCtx server.ServerCtx) error {
	if err := s.Servlets.OnStart(serverCtx); nil != err {
		return err
	}

	return nil
}

func (s *DataServlets) OnStop(serverCtx server.ServerCtx) {

	s.Servlets.OnStop(serverCtx)
}

func (s *DataServlets) Destroy(serverCtx server.ServerCtx) {

	s.Servlets.Destroy(serverCtx)
}

func (s *DataServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) {
	// probe key extraction
	bMethod := ctx.Request.Header.Peek(occp.HTTPRequestHeaderKey_Probe_Method)
	if nil == bMethod {
		return nil, fmt.Errorf("Unexpected probe method: %v", bMethod)
	}

	method := string(bMethod)

	switch method {
	case occp.HTTPRequestHeaderValue_Probe_Method_Connect:
	default:
		return nil, fmt.Errorf("Unexpected probe method: %s", method)
	}

	bProbeKey := ctx.Request.Header.Peek(occp.HTTPRequestHeaderKey_Probe_ProbeKey)
	if nil == bProbeKey {
		return nil, fmt.Errorf("Unexpected probe key : %v", bProbeKey)
	}

	probeKey := string(bProbeKey)

	grpcCTX := context.Background()
	_, err := grpc.Exec(grpcCTX, "ProbeService.readByProbeKey", probeKey)
	if nil != err {
		return nil, fmt.Errorf("grpc call Error: %s", err.Error())
	}

	servletCtx.SetAttribute(og.SessionIDKey, probeKey)
	servletCtx.SetAttribute(og.SessionClientTypeKey, og.PROBE)
	servletCtx.SetAttribute(og.SessionTargetIDKey, probeKey)

	return nil, nil
}

func (s *DataServlets) OnConnect(servletCtx server.ServletCtx, conn socket.Conn) {
	s.Servlets.OnConnect(servletCtx, conn)
}

func (s *DataServlets) OnDisconnect(servletCtx server.ServletCtx) {
	s.Servlets.OnDisconnect(servletCtx)
}

func (s *DataServlets) Handle(servletCtx server.ServletCtx,
	stopChan <-chan struct{}, doneChan chan<- struct{},
	readChan <-chan []byte, writeChan chan<- []byte) {
	defer func() {
		doneChan <- struct{}{}
	}()

	var (
		err error
	)

	for {
		select {
		case msg, ok := <-readChan:
			if !ok {
				return
			}

			err = kafka.Write("metric", msg, msg)
			if nil != err {
				logging.Logger().Error(err)
				break
			}

		case <-stopChan:
			return
		}

	}

}