add mutex

Robbie6666 · Nov 30, 2020 · b61c91e · b61c91e
1 parent 39fb64d
commit b61c91e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/pkg/healthcheck/delaycheck.go b/pkg/healthcheck/delaycheck.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/Sansui233/proxypool/pkg/proxy"
+	"sync"
 	"time"
 
 	"github.com/ivpusic/grpool"
@@ -19,6 +20,7 @@ func CleanBadProxiesWithGrpool(proxies []proxy.Proxy) (cproxies []proxy.Proxy) {
 
 	c := make(chan *Stat)
 	defer close(c)
+	m := sync.Mutex{}
 
 	pool.WaitCount(len(proxies))
 	// 线程：延迟测试，测试过程通过grpool的job并发
@@ -29,6 +31,7 @@ func CleanBadProxiesWithGrpool(proxies []proxy.Proxy) (cproxies []proxy.Proxy) {
 				defer pool.JobDone()
 				delay, err := testDelay(pp)
 				if err == nil {
+					m.Lock()
 					if ps, ok := ProxyStats.Find(p); ok {
 						ps.UpdatePSDelay(delay)
 						c <- ps
@@ -40,6 +43,7 @@ func CleanBadProxiesWithGrpool(proxies []proxy.Proxy) (cproxies []proxy.Proxy) {
 						ProxyStats = append(ProxyStats, *ps)
 						c <- ps
 					}
+					m.Unlock()
 				}
 			}
 		}

diff --git a/pkg/healthcheck/speedcheck.go b/pkg/healthcheck/speedcheck.go
@@ -33,6 +33,7 @@ func SpeedTestAll(proxies []proxy.Proxy, conns int) {
 		numJob = (numWorker + 2) / 4
 	}
 	resultCount := 0
+	m := sync.Mutex{}
 
 	log.Infoln("Speed Test ON")
 	doneCount := 0
@@ -45,6 +46,7 @@ func SpeedTestAll(proxies []proxy.Proxy, conns int) {
 			defer pool.JobDone()
 			speed, err := ProxySpeedTest(pp)
 			if err == nil || speed > 0 {
+				m.Lock()
 				if proxyStat, ok := ProxyStats.Find(pp); ok {
 					proxyStat.UpdatePSSpeed(speed)
 				} else {
@@ -54,6 +56,7 @@ func SpeedTestAll(proxies []proxy.Proxy, conns int) {
 					})
 				}
 				resultCount++
+				m.Unlock()
 			}
 			doneCount++
 			progress := float64(doneCount) * 100 / float64(len(proxies))
@@ -80,6 +83,7 @@ func SpeedTestNew(proxies []proxy.Proxy, conns int) {
 		numJob = (numWorker + 2) / 4
 	}
 	resultCount := 0
+	m := sync.Mutex{}
 
 	log.Infoln("Speed Test ON")
 	doneCount := 0
@@ -90,6 +94,7 @@ func SpeedTestNew(proxies []proxy.Proxy, conns int) {
 		pp := p
 		pool.JobQueue <- func() {
 			defer pool.JobDone()
+			m.Lock()
 			if proxyStat, ok := ProxyStats.Find(pp); !ok {
 				// when proxy's Stat not exits
 				speed, err := ProxySpeedTest(pp)
@@ -107,6 +112,7 @@ func SpeedTestNew(proxies []proxy.Proxy, conns int) {
 					resultCount++
 				}
 			}
+			m.Unlock()
 			doneCount++
 			progress := float64(doneCount) * 100 / float64(len(proxies))
 			fmt.Printf("\r\t[%5.1f%% DONE]", progress)