github/nps
1
0
Fork 0
mirror of https://github.com/ehang-io/nps.git synced 2025-07-03 04:53:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

add wait after invoke goroutine pool, fix #391

Browse Source
This commit is contained in:
ffdfgdfg 2020-02-09 00:44:09 +08:00
parent 72c695de4a
commit 632bd12bfa
2 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
lib/conn/conn.go
View File

@ -14,6 +14,7 @@ import (
"net/url" "net/url"
"strconv" "strconv"
"strings" "strings"
"sync"
"time" "time"
"ehang.io/nps/lib/common" "ehang.io/nps/lib/common"
@ -371,7 +372,10 @@ func CopyWaitGroup(conn1, conn2 net.Conn, crypt bool, snappy bool, rate *rate.Ra
//if flow != nil { //if flow != nil {
// flow.Add(in, out) // flow.Add(in, out)
//} //}
err := goroutine.CopyConnsPool.Invoke(goroutine.NewConns(connHandle, conn2, flow)) wg := new(sync.WaitGroup)
wg.Add(1)
err := goroutine.CopyConnsPool.Invoke(goroutine.NewConns(connHandle, conn2, flow, wg))
wg.Wait()
if err != nil { if err != nil {
logs.Error(err) logs.Error(err)
} }

5
lib/goroutine/pool.go
View File

@ -44,13 +44,15 @@ type Conns struct {
conn1 io.ReadWriteCloser // mux connection conn1 io.ReadWriteCloser // mux connection
conn2 net.Conn // outside connection conn2 net.Conn // outside connection
flow *file.Flow flow *file.Flow
wg *sync.WaitGroup
} }
func NewConns(c1 io.ReadWriteCloser, c2 net.Conn, flow *file.Flow) Conns { func NewConns(c1 io.ReadWriteCloser, c2 net.Conn, flow *file.Flow, wg *sync.WaitGroup) Conns {
return Conns{ return Conns{
conn1: c1, conn1: c1,
conn2: c2, conn2: c2,
flow: flow, flow: flow,
wg: wg,
} }
} }
@ -67,6 +69,7 @@ func copyConns(group interface{}) {
if conns.flow != nil { if conns.flow != nil {
conns.flow.Add(in, out) conns.flow.Add(in, out)
} }
conns.wg.Done()
} }
var connCopyPool, _ = ants.NewPoolWithFunc(200000, copyConnGroup, ants.WithNonblocking(false)) var connCopyPool, _ = ants.NewPoolWithFunc(200000, copyConnGroup, ants.WithNonblocking(false))