probe_gateway_rpc/servlet/probe-servlet.go

package servlet

import (
	"sync"
	"fmt"
	"context"
	"encoding/json"

	"github.com/valyala/fasthttp"

	"git.loafle.net/commons/server-go"
	og "git.loafle.net/overflow/gateway"
	ocpc "git.loafle.net/overflow/commons-go/probe/constants"
	ocpm "git.loafle.net/overflow/commons-go/probe/model"

	"git.loafle.net/overflow/member_gateway_rpc/subscribe"
	ogs "git.loafle.net/overflow/gateway/subscribe"
	ogrs "git.loafle.net/overflow/gateway_rpc/servlet"
	"git.loafle.net/commons/logging-go"
	"git.loafle.net/overflow/gateway/external/grpc"
	"git.loafle.net/commons/server-go/socket"

)

type ProbeServlet interface {
	ogrs.RPCServlet
}

type ProbeServlets struct {
	ogrs.RPCServlets
	connections sync.Map
}

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

	return nil
}

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

	subscribeChan, err := subscribe.Subscriber.Subscribe(ocpc.HTTPEntry_Probe)
	if nil != err {
		return err
	}
	go s.handleSubscribe(serverCtx, subscribeChan)

	return nil
}

func (s *ProbeServlets) OnStop(serverCtx server.ServerCtx) {
	if err := subscribe.Subscriber.Unsubscribe(ocpc.HTTPEntry_Probe); nil != err {
		logging.Logger().Warn(err)
	}

	s.RPCServlets.OnStop(serverCtx)
}

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

	s.RPCServlets.Destroy(serverCtx)
}

func (s *ProbeServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) {
	bMethod := ctx.Request.Header.Peek(ocpc.HTTPRequestHeaderKey_Probe_Method)

	if nil == bMethod {
		return nil, fmt.Errorf("Unexpected probe method: %v", bMethod)
	}

	method := string(bMethod)

	switch method {
	case ocpc.HTTPRequestHeaderValue_Probe_Method_Connect:
	default:
		return nil, fmt.Errorf("Unexpected noauth probe httpRequestHeaderValue: %v", method)
	}

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

	probeKey := string(bProbeKey)

	grpcCTX := context.Background()
	r, err := grpc.Exec(grpcCTX, "ProbeService.readByProbeKey", probeKey)

	if nil != err {
		return nil, fmt.Errorf("grpc call Error: %s", err.Error())
	}

	probe := ocpm.Probe{}
	err = json.Unmarshal([]byte(r), probe)
	if nil != err {
		return nil, fmt.Errorf("grpc result unMarshal Error: %s", err.Error())
	}

	extHeader := &fasthttp.ResponseHeader{}
	extHeader.Add(ocpc.HTTPResponseHeaderKey_Probe_SetEncryptionKey, probe.EncryptionKey)

	servletCtx.SetAttribute(og.SessionIDKey, probe.ProbeKey)
	servletCtx.SetAttribute(og.SessionClientTypeKey, og.PROBE)
	servletCtx.SetAttribute(og.SessionTargetIDKey, probe.ProbeKey)

	return extHeader, nil
}

func (s *ProbeServlets) OnConnect(servletCtx server.ServletCtx, conn socket.Conn) {
	s.RPCServlets.OnConnect(servletCtx, conn)

	sessionID := servletCtx.GetAttribute(og.SessionIDKey)
	targetID := servletCtx.GetAttribute(og.SessionTargetIDKey)
	if nil != sessionID && nil != targetID {
		s.connections.Store(sessionID.(string), retainConnection(targetID.(string), servletCtx))
	}
}

func (s *ProbeServlets) OnDisconnect(servletCtx server.ServletCtx) {
	s.RPCServlets.OnDisconnect(servletCtx)

	sessionID := servletCtx.GetAttribute(og.SessionIDKey)
	if nil != sessionID {
		s.connections.Delete(sessionID.(string))
	}
}

func (s *ProbeServlets) handleSubscribe(serverCtx server.ServerCtx, subscribeChan <-chan *ogs.Message) {
	for {
		select {
		case msg, ok := <- subscribeChan:
			if !ok {
				return
			}

			switch msg.TargetType {
			case ogs.PROBE:
				for _, targetID := range msg.Targets  {
					_connections := s.getProbeConnections(targetID)
					if nil == _connections || 0 == len(_connections) {
						break
					}

					for _, _connection := range _connections {
						_writeChan := _connection.servletCtx.GetAttribute(og.SessionWriteChanKey)
						if nil != _writeChan {
							writeChan := _writeChan.(chan<- []byte)
							writeChan <- msg.Message
						}
					}
				}

			}
		}
	}
}

func (s *ProbeServlets) getProbeConnections(targetID string) []*connection {
	var connections []*connection

	s.connections.Range(func(k, v interface{}) bool {
		_connection := v.(*connection)
		if _connection.targetID == targetID {
			connections = append(connections, _connection)
		}
		return true
	})

	return connections
}

//type connection struct {
//	targetID   string
//	servletCtx server.ServletCtx
//}
//
//var connectionPool sync.Pool
//
//func retainConnection(targetID string, servletCtx server.ServletCtx) *connection {
//	return nil
//}
//
//func releaseConnection(_connection *connection) {
//	_connection.targetID = ""
//	_connection.servletCtx = nil
//
//	connectionPool.Put(_connection)
//}