add wait after invoke goroutine pool, fix #391

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

View File

@ -14,6 +14,7 @@ import (
"net/url"
"strconv"
"strings"
"sync"
"time"
"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 {
// 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 {
logs.Error(err)
}

View File

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