From a15effa1ad49b22b0ba91426889a1bef01f4f1f6 Mon Sep 17 00:00:00 2001
From: crusader <crusader@loafle.com>
Date: Tue, 27 Mar 2018 00:43:44 +0900
Subject: [PATCH] ing

---
 discovery/discovery.go | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/discovery/discovery.go b/discovery/discovery.go
index 6d6144a..ad91d9b 100644
--- a/discovery/discovery.go
+++ b/discovery/discovery.go
@@ -32,6 +32,7 @@ func DiscoverZone(probeService *oopcs.ProbeService, requesterID string, dz *disc
 		return
 	}
 
+	discoverer.stopWg.Add(1)
 	discoverer.discoverZone(probeService, requesterID, dz)
 
 	discoverer.stopWg.Wait()
@@ -44,6 +45,7 @@ func DiscoverHost(probeService *oopcs.ProbeService, requesterID string, zone *di
 		logging.Logger().Warnf("Discovery: Discovery is running already")
 		return
 	}
+	discoverer.stopWg.Add(1)
 	discoverer.discoverHost(probeService, requesterID, zone, dh)
 	discoverer.stopWg.Wait()
 	discoverer.complete()
@@ -55,6 +57,7 @@ func DiscoverPort(probeService *oopcs.ProbeService, requesterID string, host *di
 		logging.Logger().Warnf("Discovery: Discovery is running already")
 		return
 	}
+	discoverer.stopWg.Add(1)
 	discoverer.discoverPort(probeService, requesterID, host, dp)
 	discoverer.stopWg.Wait()
 	discoverer.complete()
@@ -66,6 +69,7 @@ func DiscoverService(probeService *oopcs.ProbeService, requesterID string, port
 		logging.Logger().Warnf("Discovery: Discovery is running already")
 		return
 	}
+	discoverer.stopWg.Add(1)
 	discoverer.discoverService(probeService, requesterID, port, ds)
 	discoverer.stopWg.Wait()
 	discoverer.complete()
@@ -129,6 +133,9 @@ func (d *discovery) stop() {
 }
 
 func (d *discovery) discoverZone(probeService *oopcs.ProbeService, requesterID string, dz *discoveryM.DiscoveryZone) {
+	defer func() {
+		d.stopWg.Done()
+	}()
 	go taskScan(d,
 		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
 			scanZone(dz, resultChan, errChan, doneChan, stopChan)
@@ -144,6 +151,7 @@ func (d *discovery) discoverZone(probeService *oopcs.ProbeService, requesterID s
 					LastScanRange:  cr.Last().String(),
 					DiscoveryPort:  dz.DiscoveryHost.DiscoveryPort,
 				}
+				d.stopWg.Add(1)
 				d.discoverHost(probeService, requesterID, z, dh)
 			}
 		},
@@ -151,6 +159,9 @@ func (d *discovery) discoverZone(probeService *oopcs.ProbeService, requesterID s
 }
 
 func (d *discovery) discoverHost(probeService *oopcs.ProbeService, requesterID string, zone *discoveryM.Zone, dh *discoveryM.DiscoveryHost) {
+	defer func() {
+		d.stopWg.Done()
+	}()
 	go taskScan(d,
 		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
 			scanHost(zone, dh, resultChan, errChan, doneChan, stopChan)
@@ -160,6 +171,7 @@ func (d *discovery) discoverHost(probeService *oopcs.ProbeService, requesterID s
 			logging.Logger().Debugf("host: %v", h)
 			probeService.Send("DiscoveryService.DiscoveredHost", requesterID, h)
 			if nil != dh.DiscoveryPort {
+				d.stopWg.Add(1)
 				d.discoverPort(probeService, requesterID, h, dh.DiscoveryPort)
 			}
 		},
@@ -167,6 +179,9 @@ func (d *discovery) discoverHost(probeService *oopcs.ProbeService, requesterID s
 }
 
 func (d *discovery) discoverPort(probeService *oopcs.ProbeService, requesterID string, host *discoveryM.Host, dp *discoveryM.DiscoveryPort) {
+	defer func() {
+		d.stopWg.Done()
+	}()
 	go taskScan(d,
 		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
 			scanPort(host, dp, resultChan, errChan, doneChan, stopChan)
@@ -176,6 +191,7 @@ func (d *discovery) discoverPort(probeService *oopcs.ProbeService, requesterID s
 			logging.Logger().Debugf("port: %v", p)
 			probeService.Send("DiscoveryService.DiscoveredPort", requesterID, p)
 			if nil != dp.DiscoveryService {
+				d.stopWg.Add(1)
 				d.discoverService(probeService, requesterID, p, dp.DiscoveryService)
 			}
 		},
@@ -183,6 +199,9 @@ func (d *discovery) discoverPort(probeService *oopcs.ProbeService, requesterID s
 }
 
 func (d *discovery) discoverService(probeService *oopcs.ProbeService, requesterID string, port *discoveryM.Port, ds *discoveryM.DiscoveryService) {
+	defer func() {
+		d.stopWg.Done()
+	}()
 	go taskScan(d,
 		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
 			scanService(port, ds, resultChan, errChan, doneChan, stopChan)
@@ -190,7 +209,6 @@ func (d *discovery) discoverService(probeService *oopcs.ProbeService, requesterI
 		func(result interface{}) {
 			s := result.(*discoveryM.Service)
 			probeService.Send("DiscoveryService.DiscoveredService", requesterID, s)
-			logging.Logger().Debugf("service: %s(%s)[%s:%s]", s.ServiceName, s.CryptoType, port.Host.IP, port.PortNumber)
 		},
 	)
 }