ing
This commit is contained in:
parent
fdb2635edc
commit
a1ecf315c6
2
.vscode/launch.json
vendored
2
.vscode/launch.json
vendored
|
@ -9,7 +9,7 @@
|
|||
"remotePath": "",
|
||||
"port": 2345,
|
||||
"host": "127.0.0.1",
|
||||
"program": "${workspaceRoot}/main.go",
|
||||
"program": "${workspaceRoot}/cmd/main.go",
|
||||
"env": {},
|
||||
"args": [],
|
||||
"showLog": true
|
||||
|
|
11
const.go
Normal file
11
const.go
Normal file
|
@ -0,0 +1,11 @@
|
|||
package probe
|
||||
|
||||
import (
|
||||
ouc "git.loafle.net/overflow/util-go/ctx"
|
||||
)
|
||||
|
||||
const (
|
||||
SessionIDKey = ouc.CtxKey("SessionID")
|
||||
SessionTargetIDKey = ouc.CtxKey("SessionTargetID")
|
||||
SessionWriteChanKey = ouc.CtxKey("SessionWriteChan")
|
||||
)
|
|
@ -64,7 +64,7 @@ func (d *ofDiscoverer) Message() <-chan types.DiscoveryMessage {
|
|||
}
|
||||
|
||||
func (d *ofDiscoverer) SendMessage(discoveryRequest types.DiscoveryRequest, messageType types.DiscoveryMessageType, data interface{}, err error) {
|
||||
d.messageChan <- retainDiscoveryMessage(discoveryRequest, messageType, data, err)
|
||||
d.messageChan <- types.MakeDiscoveryMessage(discoveryRequest, messageType, data, err)
|
||||
}
|
||||
|
||||
func (d *ofDiscoverer) Shutdown() {
|
||||
|
@ -127,27 +127,32 @@ LOOP:
|
|||
}
|
||||
req.(*ofDiscoveryRequest).dequeue <- true
|
||||
|
||||
d.SendMessage(req, types.DiscoveryMessageTypeStart, omu.Now(), nil)
|
||||
d.discover(req)
|
||||
d.SendMessage(req, types.DiscoveryMessageTypeStop, omu.Now(), nil)
|
||||
startTime := time.Now()
|
||||
|
||||
d.SendMessage(req, types.DiscoveryMessageTypeStart, omu.Now(), nil)
|
||||
s := session.RetainDiscoverySession()
|
||||
d.discover(req, s)
|
||||
d.SendMessage(req, types.DiscoveryMessageTypeStop, time.Since(startTime).String(), nil)
|
||||
|
||||
session.ReleaseDiscoverySession(s)
|
||||
|
||||
timer := time.NewTimer(time.Millisecond * 500)
|
||||
select {
|
||||
case <-timer.C:
|
||||
}
|
||||
req.(*ofDiscoveryRequest).release()
|
||||
|
||||
case <-d.stopChan:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (d *ofDiscoverer) discover(req types.DiscoveryRequest) {
|
||||
func (d *ofDiscoverer) discover(req types.DiscoveryRequest, s session.DiscoverySession) {
|
||||
if types.DiscoveryRequestTypeNone == req.RequestType() {
|
||||
return
|
||||
}
|
||||
|
||||
s := session.RetainDiscoverySession()
|
||||
defer func() {
|
||||
session.ReleaseDiscoverySession(s)
|
||||
}()
|
||||
|
||||
if err := s.InitWithRequest(req); nil != err {
|
||||
d.SendMessage(req, types.DiscoveryMessageTypeError, nil, err)
|
||||
return
|
||||
|
|
|
@ -1,70 +0,0 @@
|
|||
package discovery
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"git.loafle.net/overflow_scanner/probe/discovery/types"
|
||||
)
|
||||
|
||||
type ofDiscoveryMessage struct {
|
||||
request types.DiscoveryRequest
|
||||
messageType types.DiscoveryMessageType
|
||||
data interface{}
|
||||
err error
|
||||
}
|
||||
|
||||
func (dm *ofDiscoveryMessage) Request() types.DiscoveryRequest {
|
||||
return dm.request
|
||||
}
|
||||
|
||||
func (dm *ofDiscoveryMessage) Type() types.DiscoveryMessageType {
|
||||
return dm.messageType
|
||||
}
|
||||
|
||||
func (dm *ofDiscoveryMessage) Data() interface{} {
|
||||
return dm.data
|
||||
}
|
||||
|
||||
func (dm *ofDiscoveryMessage) Error() error {
|
||||
return dm.err
|
||||
}
|
||||
|
||||
func (dm *ofDiscoveryMessage) Release() {
|
||||
releaseDiscoveryMessage(dm)
|
||||
}
|
||||
|
||||
func RetainDiscoveryMessage(request types.DiscoveryRequest, messageType types.DiscoveryMessageType, data interface{}, err error) *ofDiscoveryMessage {
|
||||
return retainDiscoveryMessage(request, messageType, data, err)
|
||||
}
|
||||
|
||||
func ReleaseDiscoveryMessage(dm *ofDiscoveryMessage) {
|
||||
releaseDiscoveryMessage(dm)
|
||||
}
|
||||
|
||||
var discoveryMessagePool sync.Pool
|
||||
|
||||
func retainDiscoveryMessage(request types.DiscoveryRequest, messageType types.DiscoveryMessageType, data interface{}, err error) *ofDiscoveryMessage {
|
||||
v := discoveryMessagePool.Get()
|
||||
var dm *ofDiscoveryMessage
|
||||
if v == nil {
|
||||
dm = &ofDiscoveryMessage{}
|
||||
} else {
|
||||
dm = v.(*ofDiscoveryMessage)
|
||||
}
|
||||
|
||||
dm.request = request
|
||||
dm.messageType = messageType
|
||||
dm.data = data
|
||||
dm.err = err
|
||||
|
||||
return dm
|
||||
}
|
||||
|
||||
func releaseDiscoveryMessage(dm *ofDiscoveryMessage) {
|
||||
dm.request = nil
|
||||
dm.messageType = types.DiscoveryMessageTypeNone
|
||||
dm.data = nil
|
||||
dm.err = nil
|
||||
|
||||
discoveryMessagePool.Put(dm)
|
||||
}
|
|
@ -568,6 +568,12 @@ func RetainDiscoverySession() *ofDiscoverySession {
|
|||
func ReleaseDiscoverySession(ds *ofDiscoverySession) {
|
||||
close(ds.stopChan)
|
||||
|
||||
// timer := time.NewTimer(time.Microsecond * 500)
|
||||
|
||||
// select {
|
||||
// case <-timer.C:
|
||||
// }
|
||||
|
||||
ds.discoveryRequest = nil
|
||||
ds.zone = nil
|
||||
ds.host = nil
|
||||
|
|
|
@ -1,40 +1,5 @@
|
|||
package types
|
||||
|
||||
func NewMockDiscoveryMessage(request DiscoveryRequest, messageType DiscoveryMessageType, data interface{}, err error) DiscoveryMessage {
|
||||
return &MockDiscoveryMessage{
|
||||
request: request,
|
||||
messageType: messageType,
|
||||
data: data,
|
||||
err: err,
|
||||
}
|
||||
}
|
||||
|
||||
type MockDiscoveryMessage struct {
|
||||
request DiscoveryRequest
|
||||
messageType DiscoveryMessageType
|
||||
data interface{}
|
||||
err error
|
||||
}
|
||||
|
||||
func (dm *MockDiscoveryMessage) Request() DiscoveryRequest {
|
||||
return dm.request
|
||||
}
|
||||
|
||||
func (dm *MockDiscoveryMessage) Type() DiscoveryMessageType {
|
||||
return dm.messageType
|
||||
}
|
||||
|
||||
func (dm *MockDiscoveryMessage) Data() interface{} {
|
||||
return dm.data
|
||||
}
|
||||
|
||||
func (dm *MockDiscoveryMessage) Error() error {
|
||||
return dm.err
|
||||
}
|
||||
|
||||
func (dm *MockDiscoveryMessage) Release() {
|
||||
}
|
||||
|
||||
func NewMockDiscoveryRequest(requesterID string, requestType DiscoveryRequestType, params ...interface{}) DiscoveryRequest {
|
||||
return &MockDiscoveryRequest{
|
||||
requesterID: requesterID,
|
||||
|
|
|
@ -15,14 +15,6 @@ const (
|
|||
DiscoveryMessageTypeService
|
||||
)
|
||||
|
||||
type DiscoveryMessage interface {
|
||||
Request() DiscoveryRequest
|
||||
Type() DiscoveryMessageType
|
||||
Data() interface{}
|
||||
Error() error
|
||||
Release()
|
||||
}
|
||||
|
||||
type DiscoveryRequestType int
|
||||
|
||||
const (
|
||||
|
@ -32,6 +24,14 @@ const (
|
|||
DiscoveryRequestTypeService
|
||||
)
|
||||
|
||||
type DiscoveryMessage func() (request DiscoveryRequest, messageType DiscoveryMessageType, data interface{}, err error)
|
||||
|
||||
func MakeDiscoveryMessage(request DiscoveryRequest, messageType DiscoveryMessageType, data interface{}, err error) DiscoveryMessage {
|
||||
return func() (DiscoveryRequest, DiscoveryMessageType, interface{}, error) {
|
||||
return request, messageType, data, err
|
||||
}
|
||||
}
|
||||
|
||||
type DiscoveryRequest interface {
|
||||
RequesterID() string
|
||||
RequestType() DiscoveryRequestType
|
||||
|
|
223
internal/pubsub/pubsub.go
Normal file
223
internal/pubsub/pubsub.go
Normal file
|
@ -0,0 +1,223 @@
|
|||
package pubsub
|
||||
|
||||
type operation int
|
||||
|
||||
const (
|
||||
sub operation = iota
|
||||
subOnce
|
||||
pub
|
||||
tryPub
|
||||
unsub
|
||||
unsubAll
|
||||
closeTopic
|
||||
shutdown
|
||||
)
|
||||
|
||||
// PubSub is a collection of topics.
|
||||
type PubSub struct {
|
||||
cmdChan chan cmd
|
||||
capacity int
|
||||
}
|
||||
|
||||
type cmd struct {
|
||||
op operation
|
||||
topics []string
|
||||
ch chan interface{}
|
||||
msg interface{}
|
||||
}
|
||||
|
||||
// New creates a new PubSub and starts a goroutine for handling operations.
|
||||
// The capacity of the channels created by Sub and SubOnce will be as specified.
|
||||
func New(capacity int) *PubSub {
|
||||
ps := &PubSub{make(chan cmd), capacity}
|
||||
go ps.start()
|
||||
return ps
|
||||
}
|
||||
|
||||
// Sub returns a channel on which messages published on any of
|
||||
// the specified topics can be received.
|
||||
func (ps *PubSub) Sub(topics ...string) chan interface{} {
|
||||
return ps.sub(sub, topics...)
|
||||
}
|
||||
|
||||
// SubOnce is similar to Sub, but only the first message published, after subscription,
|
||||
// on any of the specified topics can be received.
|
||||
func (ps *PubSub) SubOnce(topics ...string) chan interface{} {
|
||||
return ps.sub(subOnce, topics...)
|
||||
}
|
||||
|
||||
func (ps *PubSub) sub(op operation, topics ...string) chan interface{} {
|
||||
ch := make(chan interface{}, ps.capacity)
|
||||
ps.cmdChan <- cmd{op: op, topics: topics, ch: ch}
|
||||
return ch
|
||||
}
|
||||
|
||||
// AddSub adds subscriptions to an existing channel.
|
||||
func (ps *PubSub) AddSub(ch chan interface{}, topics ...string) {
|
||||
ps.cmdChan <- cmd{op: sub, topics: topics, ch: ch}
|
||||
}
|
||||
|
||||
// Pub publishes the given message to all subscribers of
|
||||
// the specified topics.
|
||||
func (ps *PubSub) Pub(msg interface{}, topics ...string) {
|
||||
ps.cmdChan <- cmd{op: pub, topics: topics, msg: msg}
|
||||
}
|
||||
|
||||
// TryPub publishes the given message to all subscribers of
|
||||
// the specified topics if the topic has buffer space.
|
||||
func (ps *PubSub) TryPub(msg interface{}, topics ...string) {
|
||||
ps.cmdChan <- cmd{op: tryPub, topics: topics, msg: msg}
|
||||
}
|
||||
|
||||
// Unsub unsubscribes the given channel from the specified
|
||||
// topics. If no topic is specified, it is unsubscribed
|
||||
// from all topics.
|
||||
func (ps *PubSub) Unsub(ch chan interface{}, topics ...string) {
|
||||
if len(topics) == 0 {
|
||||
ps.cmdChan <- cmd{op: unsubAll, ch: ch}
|
||||
return
|
||||
}
|
||||
|
||||
ps.cmdChan <- cmd{op: unsub, topics: topics, ch: ch}
|
||||
}
|
||||
|
||||
// Close closes all channels currently subscribed to the specified topics.
|
||||
// If a channel is subscribed to multiple topics, some of which is
|
||||
// not specified, it is not closed.
|
||||
func (ps *PubSub) Close(topics ...string) {
|
||||
ps.cmdChan <- cmd{op: closeTopic, topics: topics}
|
||||
}
|
||||
|
||||
// Shutdown closes all subscribed channels and terminates the goroutine.
|
||||
func (ps *PubSub) Shutdown() {
|
||||
ps.cmdChan <- cmd{op: shutdown}
|
||||
}
|
||||
|
||||
func (ps *PubSub) start() {
|
||||
reg := registry{
|
||||
topics: make(map[string]map[chan interface{}]bool),
|
||||
revTopics: make(map[chan interface{}]map[string]bool),
|
||||
}
|
||||
|
||||
loop:
|
||||
for cmd := range ps.cmdChan {
|
||||
if cmd.topics == nil {
|
||||
switch cmd.op {
|
||||
case unsubAll:
|
||||
reg.removeChannel(cmd.ch)
|
||||
|
||||
case shutdown:
|
||||
break loop
|
||||
}
|
||||
|
||||
continue loop
|
||||
}
|
||||
|
||||
for _, topic := range cmd.topics {
|
||||
switch cmd.op {
|
||||
case sub:
|
||||
reg.add(topic, cmd.ch, false)
|
||||
|
||||
case subOnce:
|
||||
reg.add(topic, cmd.ch, true)
|
||||
|
||||
case tryPub:
|
||||
reg.sendNoWait(topic, cmd.msg)
|
||||
|
||||
case pub:
|
||||
reg.send(topic, cmd.msg)
|
||||
|
||||
case unsub:
|
||||
reg.remove(topic, cmd.ch)
|
||||
|
||||
case closeTopic:
|
||||
reg.removeTopic(topic)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for topic, chans := range reg.topics {
|
||||
for ch := range chans {
|
||||
reg.remove(topic, ch)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// registry maintains the current subscription state. It's not
|
||||
// safe to access a registry from multiple goroutines simultaneously.
|
||||
type registry struct {
|
||||
topics map[string]map[chan interface{}]bool
|
||||
revTopics map[chan interface{}]map[string]bool
|
||||
}
|
||||
|
||||
func (reg *registry) add(topic string, ch chan interface{}, once bool) {
|
||||
if reg.topics[topic] == nil {
|
||||
reg.topics[topic] = make(map[chan interface{}]bool)
|
||||
}
|
||||
reg.topics[topic][ch] = once
|
||||
|
||||
if reg.revTopics[ch] == nil {
|
||||
reg.revTopics[ch] = make(map[string]bool)
|
||||
}
|
||||
reg.revTopics[ch][topic] = true
|
||||
}
|
||||
|
||||
func (reg *registry) send(topic string, msg interface{}) {
|
||||
for ch, once := range reg.topics[topic] {
|
||||
ch <- msg
|
||||
if once {
|
||||
for topic := range reg.revTopics[ch] {
|
||||
reg.remove(topic, ch)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (reg *registry) sendNoWait(topic string, msg interface{}) {
|
||||
for ch, once := range reg.topics[topic] {
|
||||
select {
|
||||
case ch <- msg:
|
||||
if once {
|
||||
for topic := range reg.revTopics[ch] {
|
||||
reg.remove(topic, ch)
|
||||
}
|
||||
}
|
||||
default:
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
func (reg *registry) removeTopic(topic string) {
|
||||
for ch := range reg.topics[topic] {
|
||||
reg.remove(topic, ch)
|
||||
}
|
||||
}
|
||||
|
||||
func (reg *registry) removeChannel(ch chan interface{}) {
|
||||
for topic := range reg.revTopics[ch] {
|
||||
reg.remove(topic, ch)
|
||||
}
|
||||
}
|
||||
|
||||
func (reg *registry) remove(topic string, ch chan interface{}) {
|
||||
if _, ok := reg.topics[topic]; !ok {
|
||||
return
|
||||
}
|
||||
|
||||
if _, ok := reg.topics[topic][ch]; !ok {
|
||||
return
|
||||
}
|
||||
|
||||
delete(reg.topics[topic], ch)
|
||||
delete(reg.revTopics[ch], topic)
|
||||
|
||||
if len(reg.topics[topic]) == 0 {
|
||||
delete(reg.topics, topic)
|
||||
}
|
||||
|
||||
if len(reg.revTopics[ch]) == 0 {
|
||||
close(ch)
|
||||
delete(reg.revTopics, ch)
|
||||
}
|
||||
}
|
7
internal/rpc/message.go
Normal file
7
internal/rpc/message.go
Normal file
|
@ -0,0 +1,7 @@
|
|||
package rpc
|
||||
|
||||
type RPCMessage func() ([]string, string, []interface{})
|
||||
|
||||
func MakeRPCMessage(targets []string, method string, params []interface{}) RPCMessage {
|
||||
return func() ([]string, string, []interface{}) { return targets, method, params }
|
||||
}
|
|
@ -1,12 +1,21 @@
|
|||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"sync"
|
||||
|
||||
"git.loafle.net/overflow_scanner/probe/internal/pubsub"
|
||||
"git.loafle.net/overflow_scanner/probe/internal/rpc"
|
||||
|
||||
olog "git.loafle.net/overflow/log-go"
|
||||
orp "git.loafle.net/overflow/rpc-go/protocol"
|
||||
orr "git.loafle.net/overflow/rpc-go/registry"
|
||||
"git.loafle.net/overflow/server-go"
|
||||
oss "git.loafle.net/overflow/server-go/socket"
|
||||
ossw "git.loafle.net/overflow/server-go/socket/web"
|
||||
"git.loafle.net/overflow_scanner/probe"
|
||||
uuid "github.com/satori/go.uuid"
|
||||
"github.com/valyala/fasthttp"
|
||||
)
|
||||
|
||||
|
@ -19,6 +28,11 @@ type ScannerServlets struct {
|
|||
|
||||
RPCInvoker orr.RPCInvoker
|
||||
RPCServerCodec orp.ServerCodec
|
||||
PubSub *pubsub.PubSub
|
||||
|
||||
sessions sync.Map
|
||||
|
||||
subscribeChan chan interface{}
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) Init(serverCtx server.ServerCtx) error {
|
||||
|
@ -34,11 +48,16 @@ func (s *ScannerServlets) OnStart(serverCtx server.ServerCtx) error {
|
|||
return err
|
||||
}
|
||||
|
||||
s.subscribeChan = s.PubSub.Sub("/scanner")
|
||||
go s.handleSubscribe(serverCtx)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) OnStop(serverCtx server.ServerCtx) {
|
||||
|
||||
s.PubSub.Unsub(s.subscribeChan, "/scanner")
|
||||
|
||||
s.Servlets.OnStop(serverCtx)
|
||||
}
|
||||
|
||||
|
@ -48,6 +67,15 @@ func (s *ScannerServlets) Destroy(serverCtx server.ServerCtx) {
|
|||
}
|
||||
|
||||
func (s *ScannerServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) {
|
||||
requesterID := string(ctx.QueryArgs().Peek("requesterID"))
|
||||
if "" == requesterID {
|
||||
return nil, fmt.Errorf("requesterID is not valid")
|
||||
}
|
||||
|
||||
sessionID := uuid.NewV4().String()
|
||||
|
||||
servletCtx.SetAttribute(probe.SessionTargetIDKey, requesterID)
|
||||
servletCtx.SetAttribute(probe.SessionIDKey, sessionID)
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
@ -55,11 +83,25 @@ func (s *ScannerServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.
|
|||
func (s *ScannerServlets) OnConnect(servletCtx server.ServletCtx, conn oss.Conn) {
|
||||
s.Servlets.OnConnect(servletCtx, conn)
|
||||
|
||||
sessionID := servletCtx.GetAttribute(probe.SessionIDKey)
|
||||
targetID := servletCtx.GetAttribute(probe.SessionTargetIDKey)
|
||||
|
||||
if nil != sessionID && nil != targetID {
|
||||
s.sessions.Store(sessionID.(string), RetainSession(targetID.(string), servletCtx))
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) OnDisconnect(servletCtx server.ServletCtx) {
|
||||
s.Servlets.OnDisconnect(servletCtx)
|
||||
|
||||
sessionID := servletCtx.GetAttribute(probe.SessionIDKey)
|
||||
if nil != sessionID {
|
||||
session, ok := s.sessions.Load(sessionID)
|
||||
if ok {
|
||||
ReleaseSession(session.(*Session))
|
||||
}
|
||||
s.sessions.Delete(sessionID.(string))
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) Handle(servletCtx server.ServletCtx,
|
||||
|
@ -76,6 +118,8 @@ func (s *ScannerServlets) Handle(servletCtx server.ServletCtx,
|
|||
err error
|
||||
)
|
||||
|
||||
servletCtx.SetAttribute(probe.SessionWriteChanKey, writeChan)
|
||||
|
||||
for {
|
||||
select {
|
||||
case socketMessage, ok := <-readChan:
|
||||
|
@ -118,6 +162,48 @@ func (s *ScannerServlets) Handle(servletCtx server.ServletCtx,
|
|||
}
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) handleSubscribe(serverCtx server.ServerCtx) {
|
||||
var sessions []*Session
|
||||
|
||||
LOOP:
|
||||
for {
|
||||
select {
|
||||
case msg, ok := <-s.subscribeChan:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
_msg, ok := msg.(rpc.RPCMessage)
|
||||
if !ok {
|
||||
log.Print("RPCMessage is not valid")
|
||||
continue LOOP
|
||||
}
|
||||
|
||||
targets, method, params := _msg()
|
||||
|
||||
sessions = s.GetSessionsByTargetIDs(targets)
|
||||
|
||||
if nil == sessions || 0 == len(sessions) {
|
||||
continue LOOP
|
||||
}
|
||||
|
||||
messageType, message, err := s.RPCServerCodec.NewNotification(method, params)
|
||||
if nil != err {
|
||||
log.Print("RPCMessage is not valid ", _msg, err)
|
||||
continue LOOP
|
||||
}
|
||||
|
||||
for _, session := range sessions {
|
||||
_writeChan := session.ServletCtx.GetAttribute(probe.SessionWriteChanKey)
|
||||
if nil != _writeChan {
|
||||
writeChan := _writeChan.(chan<- oss.SocketMessage)
|
||||
writeChan <- oss.MakeSocketMessage(messageType, message)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) writeError(src orp.ServerRequestCodec, writeChan chan<- oss.SocketMessage, code orp.ErrorCode, message string, data interface{}) {
|
||||
if !src.HasResponse() {
|
||||
return
|
||||
|
@ -136,3 +222,42 @@ func (s *ScannerServlets) writeError(src orp.ServerRequestCodec, writeChan chan<
|
|||
}
|
||||
writeChan <- oss.MakeSocketMessage(resMessageType, resMessage)
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) GetSessions(sessionIDs []string) []*Session {
|
||||
var sessions []*Session
|
||||
|
||||
if nil == sessionIDs || 0 == len(sessionIDs) {
|
||||
return sessions
|
||||
}
|
||||
|
||||
for _, sessionID := range sessionIDs {
|
||||
session, ok := s.sessions.Load(sessionID)
|
||||
if ok {
|
||||
sessions = append(sessions, session.(*Session))
|
||||
}
|
||||
}
|
||||
|
||||
return sessions
|
||||
}
|
||||
|
||||
func (s *ScannerServlets) GetSessionsByTargetIDs(targetIDs []string) []*Session {
|
||||
var sessions []*Session
|
||||
if nil == targetIDs || 0 == len(targetIDs) {
|
||||
return sessions
|
||||
}
|
||||
|
||||
s.sessions.Range(func(k, v interface{}) bool {
|
||||
session := v.(*Session)
|
||||
|
||||
for _, targetID := range targetIDs {
|
||||
if session.TargetID == targetID {
|
||||
sessions = append(sessions, session)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
})
|
||||
|
||||
return sessions
|
||||
}
|
||||
|
|
|
@ -10,6 +10,7 @@ import (
|
|||
olog "git.loafle.net/overflow/log-go"
|
||||
"git.loafle.net/overflow/server-go"
|
||||
ossw "git.loafle.net/overflow/server-go/socket/web"
|
||||
"git.loafle.net/overflow_scanner/probe/internal/pubsub"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
|
@ -20,8 +21,9 @@ type ServerHandler interface {
|
|||
type ServerHandlers struct {
|
||||
ossw.ServerHandlers
|
||||
|
||||
Services []interface{}
|
||||
OrderedServices []reflect.Type
|
||||
Services []interface{}
|
||||
|
||||
PubSub *pubsub.PubSub
|
||||
}
|
||||
|
||||
func (sh *ServerHandlers) Init(serverCtx server.ServerCtx) error {
|
||||
|
@ -51,6 +53,8 @@ func (sh *ServerHandlers) OnStop(serverCtx server.ServerCtx) {
|
|||
}
|
||||
}
|
||||
|
||||
sh.PubSub.Shutdown()
|
||||
|
||||
sh.ServerHandlers.OnStop(serverCtx)
|
||||
}
|
||||
|
||||
|
@ -82,9 +86,5 @@ func (sh *ServerHandlers) Validate() error {
|
|||
return fmt.Errorf("Services must be specified")
|
||||
}
|
||||
|
||||
if nil == sh.OrderedServices {
|
||||
return fmt.Errorf("OrderedServices must be specified")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
|
|
@ -9,6 +9,7 @@ import (
|
|||
|
||||
"git.loafle.net/overflow_scanner/probe/annotation"
|
||||
"git.loafle.net/overflow_scanner/probe/discovery"
|
||||
"git.loafle.net/overflow_scanner/probe/internal/pubsub"
|
||||
|
||||
// For service
|
||||
_ "git.loafle.net/overflow_scanner/probe/service"
|
||||
|
@ -16,8 +17,10 @@ import (
|
|||
|
||||
func New() *ossw.Server {
|
||||
rpcServerCodec := orpj.NewServerCodec()
|
||||
_pubsub := pubsub.New(10)
|
||||
|
||||
od.RegisterSingletonByName("Discoverer", discovery.Instance())
|
||||
od.RegisterSingletonByName("PubSub", _pubsub)
|
||||
|
||||
services, err := od.GetInstancesByAnnotationType(annotation.ServiceAnnotationType)
|
||||
if nil != err {
|
||||
|
@ -30,10 +33,12 @@ func New() *ossw.Server {
|
|||
ss := &ScannerServlets{}
|
||||
ss.RPCInvoker = rpcRegistry
|
||||
ss.RPCServerCodec = rpcServerCodec
|
||||
ss.PubSub = _pubsub
|
||||
|
||||
sh := &ServerHandlers{}
|
||||
sh.Name = "Probe"
|
||||
sh.Services = services
|
||||
sh.PubSub = _pubsub
|
||||
|
||||
sh.RegisterServlet("/scanner", ss)
|
||||
|
||||
|
|
36
server/session.go
Normal file
36
server/session.go
Normal file
|
@ -0,0 +1,36 @@
|
|||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
server "git.loafle.net/overflow/server-go"
|
||||
)
|
||||
|
||||
type Session struct {
|
||||
TargetID string
|
||||
ServletCtx server.ServletCtx
|
||||
}
|
||||
|
||||
var sessionPool sync.Pool
|
||||
|
||||
func RetainSession(targetID string, servletCtx server.ServletCtx) *Session {
|
||||
v := sessionPool.Get()
|
||||
var _session *Session
|
||||
if v == nil {
|
||||
_session = &Session{}
|
||||
} else {
|
||||
_session = v.(*Session)
|
||||
}
|
||||
|
||||
_session.TargetID = targetID
|
||||
_session.ServletCtx = servletCtx
|
||||
|
||||
return _session
|
||||
}
|
||||
|
||||
func ReleaseSession(_session *Session) {
|
||||
_session.TargetID = ""
|
||||
_session.ServletCtx = nil
|
||||
|
||||
sessionPool.Put(_session)
|
||||
}
|
|
@ -4,6 +4,9 @@ import (
|
|||
"log"
|
||||
"reflect"
|
||||
|
||||
"git.loafle.net/overflow_scanner/probe/internal/pubsub"
|
||||
"git.loafle.net/overflow_scanner/probe/internal/rpc"
|
||||
|
||||
oa "git.loafle.net/overflow/annotation-go"
|
||||
od "git.loafle.net/overflow/di-go"
|
||||
omd "git.loafle.net/overflow/model/discovery"
|
||||
|
@ -21,6 +24,7 @@ type DiscoveryService struct {
|
|||
oa.TypeAnnotation `annotation:"@Injectable('name': 'DiscoveryService') @Service()"`
|
||||
|
||||
Discoverer discovery.Discoverer `annotation:"@Inject('name': 'Discoverer')"`
|
||||
PubSub *pubsub.PubSub `annotation:"@Inject('name': 'PubSub')"`
|
||||
|
||||
_InitService oa.MethodAnnotation `annotation:"@PostConstruct()"`
|
||||
_DestroyService oa.MethodAnnotation `annotation:"@PreDestroy()"`
|
||||
|
@ -35,26 +39,99 @@ func (s *DiscoveryService) InitService() {
|
|||
return
|
||||
}
|
||||
|
||||
switch msg.Type() {
|
||||
request, messageType, data, err := msg()
|
||||
|
||||
switch messageType {
|
||||
case types.DiscoveryMessageTypeStart:
|
||||
log.Print("Start ", msg.Data())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.DiscoveryStart",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Start ", data)
|
||||
case types.DiscoveryMessageTypeStop:
|
||||
log.Print("Stop ", msg.Data())
|
||||
return
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.DiscoveryStop",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Stop ", data)
|
||||
case types.DiscoveryMessageTypeQueueing:
|
||||
log.Print("Queueing ", msg.Data())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.Queueing",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Queueing ", data)
|
||||
case types.DiscoveryMessageTypeFailedQueueing:
|
||||
log.Print("FailedQueueing ", msg.Data())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.FailedQueueing",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("FailedQueueing ", data)
|
||||
case types.DiscoveryMessageTypeTimeout:
|
||||
log.Print("Timeout ", msg.Data())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.Timeout",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Timeout ", data)
|
||||
case types.DiscoveryMessageTypeError:
|
||||
log.Print("Error ", msg.Error())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.Error",
|
||||
[]interface{}{err},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Error ", err)
|
||||
case types.DiscoveryMessageTypeHost:
|
||||
log.Print("Host ", msg.Data())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.DiscoveredHost",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Host ", data)
|
||||
case types.DiscoveryMessageTypePort:
|
||||
log.Print("Port ", msg.Data())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.DiscoveredPort",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Port ", data)
|
||||
case types.DiscoveryMessageTypeService:
|
||||
log.Print("Service ", msg.Data())
|
||||
s.PubSub.Pub(
|
||||
rpc.MakeRPCMessage(
|
||||
[]string{request.RequesterID()},
|
||||
"DiscoveryService.DiscoveredService",
|
||||
[]interface{}{data},
|
||||
),
|
||||
"/scanner",
|
||||
)
|
||||
log.Print("Service ", data)
|
||||
default:
|
||||
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue
Block a user