From 348cb63eac4505937379dbd4c5169df2e1a33595 Mon Sep 17 00:00:00 2001
From: crusader <crusader@loafle.com>
Date: Tue, 27 Mar 2018 00:33:01 +0900
Subject: [PATCH] ing

---
 discovery/discovery.go | 52 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/discovery/discovery.go b/discovery/discovery.go
index ca2cad2..b09d90a 100644
--- a/discovery/discovery.go
+++ b/discovery/discovery.go
@@ -15,7 +15,9 @@ var RPCServlet rpc.Servlet
 var discoverer *discovery
 
 func DiscoveryInit() {
-	discoverer = &discovery{}
+	discoverer = &discovery{
+		_isRunning: false,
+	}
 	discoverer.start()
 }
 
@@ -25,19 +27,45 @@ func DiscoveryDestroy() {
 }
 
 func DiscoverZone(probeService *oopcs.ProbeService, requesterID string, dz *discoveryM.DiscoveryZone) {
+	if !discoverer.canRunOrNot() {
+		logging.Logger().Warnf("Discovery: Discovery is running already")
+		return
+	}
+
 	discoverer.discoverZone(probeService, requesterID, dz)
+
+	discoverer.stopWg.Wait()
+	discoverer.complete()
 }
 
 func DiscoverHost(probeService *oopcs.ProbeService, requesterID string, zone *discoveryM.Zone, dh *discoveryM.DiscoveryHost) {
+	if !discoverer.canRunOrNot() {
+		logging.Logger().Warnf("Discovery: Discovery is running already")
+		return
+	}
 	discoverer.discoverHost(probeService, requesterID, zone, dh)
+	discoverer.stopWg.Wait()
+	discoverer.complete()
 }
 
 func DiscoverPort(probeService *oopcs.ProbeService, requesterID string, host *discoveryM.Host, dp *discoveryM.DiscoveryPort) {
+	if !discoverer.canRunOrNot() {
+		logging.Logger().Warnf("Discovery: Discovery is running already")
+		return
+	}
 	discoverer.discoverPort(probeService, requesterID, host, dp)
+	discoverer.stopWg.Wait()
+	discoverer.complete()
 }
 
 func DiscoverService(probeService *oopcs.ProbeService, requesterID string, port *discoveryM.Port, ds *discoveryM.DiscoveryService) {
+	if !discoverer.canRunOrNot() {
+		logging.Logger().Warnf("Discovery: Discovery is running already")
+		return
+	}
 	discoverer.discoverService(probeService, requesterID, port, ds)
+	discoverer.stopWg.Wait()
+	discoverer.complete()
 }
 
 func Stop() {
@@ -50,6 +78,28 @@ type discovery struct {
 
 	stopChan chan struct{}
 	stopWg   sync.WaitGroup
+
+	_isRunning bool
+	_statusMtx sync.Mutex
+}
+
+func (d *discovery) canRunOrNot() bool {
+	d._statusMtx.Lock()
+	defer d._statusMtx.Unlock()
+
+	if d._isRunning {
+		return false
+	}
+
+	d._isRunning = true
+	return true
+}
+
+func (d *discovery) complete() {
+	d._statusMtx.Lock()
+	defer d._statusMtx.Unlock()
+
+	d._isRunning = false
 }
 
 func (d *discovery) start() {