ing

discovery/discoverer.go

@@ -44,8 +44,8 @@ type ofDiscoverer struct {
 	stopChan chan struct{}
 	stopWg   sync.WaitGroup
 
-	requestIDs         map[string]bool
-	processingSessions map[string]session.DiscoverySession
+	requestIDs         sync.Map //map[string]bool
+	processingSessions sync.Map //map[string]session.DiscoverySession
 	requestQueue       chan types.DiscoveryRequest
 	messageChan        chan types.DiscoveryMessage
 }
@@ -63,18 +63,18 @@ func (d *ofDiscoverer) DiscoverService(requesterID string, port *omd.Port, ds *o
 }
 
 func (d *ofDiscoverer) DiscoverStop(requesterID string, requestID string) {
-	_, ok := d.requestIDs[requestID]
+	_, ok := d.requestIDs.Load(requestID)
 	if ok {
-		d.requestIDs[requestID] = false
+		d.requestIDs.Store(requestID, true)
 		return
 	}
 
-	s, ok := d.processingSessions[requestID]
+	s, ok := d.processingSessions.Load(requestID)
 	if !ok {
 		return
 	}
-	delete(d.processingSessions, requestID)
-	s.Shutdown()
+	s.(session.DiscoverySession).DiscoveryRequest().(*ofDiscoveryRequest).canceled.Store(true)
+	s.(session.DiscoverySession).Shutdown()
 }
 
 func (d *ofDiscoverer) Message() <-chan types.DiscoveryMessage {
@@ -82,6 +82,10 @@ func (d *ofDiscoverer) Message() <-chan types.DiscoveryMessage {
 }
 
 func (d *ofDiscoverer) SendMessage(discoveryRequest types.DiscoveryRequest, messageType types.DiscoveryMessageType, datas ...interface{}) {
+	if discoveryRequest.(*ofDiscoveryRequest).canceled.Load().(bool) {
+		return
+	}
+
 	d.messageChan <- types.MakeDiscoveryMessage(discoveryRequest, messageType, datas...)
 }
 
@@ -102,8 +106,6 @@ func (d *ofDiscoverer) Shutdown() {
 func (d *ofDiscoverer) start() {
 	d.stopChan = make(chan struct{})
 
-	d.requestIDs = make(map[string]bool, 10)
-	d.processingSessions = make(map[string]session.DiscoverySession, 2)
 	d.requestQueue = make(chan types.DiscoveryRequest, 10)
 	d.messageChan = make(chan types.DiscoveryMessage, 256)
 
@@ -114,7 +116,7 @@ func (d *ofDiscoverer) start() {
 func (d *ofDiscoverer) enqueue(req *ofDiscoveryRequest) {
 	select {
 	case d.requestQueue <- req:
-		d.requestIDs[req.RequestID()] = true
+		d.requestIDs.Store(req.RequestID(), false)
 		d.SendMessage(req, types.DiscoveryMessageTypeQueueing, req.RequestID(), omu.Now())
 		go func() {
 			select {
@@ -147,9 +149,9 @@ LOOP:
 				return
 			}
 
-			canceled, ok := d.requestIDs[req.RequestID()]
-			delete(d.requestIDs, req.RequestID())
-			if !ok || !canceled {
+			canceled, ok := d.requestIDs.Load(req.RequestID())
+			d.requestIDs.Delete(req.RequestID())
+			if !ok || canceled.(bool) {
 				req.(*ofDiscoveryRequest).release()
 				continue LOOP
 			}
@@ -162,15 +164,15 @@ LOOP:
 
 			d.SendMessage(req, types.DiscoveryMessageTypeStart, omu.Now())
 			s := session.RetainDiscoverySession()
-			d.processingSessions[req.RequestID()] = s
+			d.processingSessions.Store(req.RequestID(), s)
 			d.discover(req, s)
 			select {
 			case <-time.After(time.Second * 1):
 			}
 			d.SendMessage(req, types.DiscoveryMessageTypeStop, omu.Now())
 
-			if _, ok := d.processingSessions[req.RequestID()]; ok {
-				delete(d.processingSessions, req.RequestID())
+			if _, ok := d.processingSessions.Load(req.RequestID()); ok {
+				d.processingSessions.Delete(req.RequestID())
 				s.Shutdown()
 			}
 

discovery/discovery-request.go

@@ -2,6 +2,7 @@ package discovery
 
 import (
 	"sync"
+	"sync/atomic"
 
 	"git.loafle.net/overflow_scanner/probe/discovery/types"
 	uuid "github.com/satori/go.uuid"
@@ -13,6 +14,8 @@ type ofDiscoveryRequest struct {
 	requestType types.DiscoveryRequestType
 	params      []interface{}
 
+	canceled atomic.Value
+
 	dequeue chan bool
 	timeout bool
 }
@@ -52,6 +55,7 @@ func retainDiscoveryRequest(requesterID string, requestType types.DiscoveryReque
 	dr.requesterID = requesterID
 	dr.requestType = requestType
 	dr.params = params
+	dr.canceled.Store(false)
 	dr.dequeue = make(chan bool, 1)
 	dr.timeout = false
 

discovery/session/discovery-session.go

@@ -559,6 +559,10 @@ func (ds *ofDiscoverySession) delegate(data interface{}) {
 }
 
 func (ds *ofDiscoverySession) Shutdown() {
+	if ds.stopped.Load().(bool) {
+		return
+	}
+
 	ds.stopped.Store(true)
 	close(ds.stopChan)
 }