diff --git a/client/highway.go b/client/highway.go index 9aa6230e..e409e065 100644 --- a/client/highway.go +++ b/client/highway.go @@ -15,6 +15,7 @@ import ( "time" "github.com/pkg/errors" + "golang.org/x/sync/errgroup" "google.golang.org/protobuf/proto" "github.com/Mrs4s/MiraiGo/binary" @@ -232,7 +233,6 @@ func (c *QQClient) highwayUploadFileMultiThreadingByBDH(path string, cmdId int32 rspExt []byte BlockId = ^uint32(0) // -1 uploadedCount uint32 - lastErr error cond = sync.NewCond(&sync.Mutex{}) ) // Init Blocks @@ -253,6 +253,8 @@ func (c *QQClient) highwayUploadFileMultiThreadingByBDH(path string, cmdId int32 }) } doUpload := func() error { + defer cond.Signal() + conn, err := net.DialTimeout("tcp", c.srvSsoAddrs[0], time.Second*20) if err != nil { return errors.Wrap(err, "connect error") @@ -279,13 +281,10 @@ func (c *QQClient) highwayUploadFileMultiThreadingByBDH(path string, cmdId int32 block := blocks[nextId] if block.Id == len(blocks)-1 { cond.L.Lock() - for atomic.LoadUint32(&uploadedCount) != uint32(len(blocks)-1) && lastErr == nil { + for atomic.LoadUint32(&uploadedCount) != uint32(len(blocks))-1 { cond.Wait() } cond.L.Unlock() - if lastErr != nil { - break - } } buffer = buffer[:blockSize] _, _ = chunk.Seek(block.BeginOffset, io.SeekStart) @@ -347,19 +346,13 @@ func (c *QQClient) highwayUploadFileMultiThreadingByBDH(path string, cmdId int32 } return nil } - wg := sync.WaitGroup{} - wg.Add(threadCount) + + group := errgroup.Group{} for i := 0; i < threadCount; i++ { - go func() { - defer wg.Done() - defer cond.Signal() - if err := doUpload(); err != nil { - lastErr = err - } - }() + group.Go(doUpload) } - wg.Wait() - return rspExt, lastErr + err = group.Wait() + return rspExt, err } func (c *QQClient) highwaySendHeartbreak(conn net.Conn) error { diff --git a/go.mod b/go.mod index 419deca3..f5b2f624 100644 --- a/go.mod +++ b/go.mod @@ -6,5 +6,6 @@ require ( github.com/json-iterator/go v1.1.10 github.com/modern-go/reflect2 v1.0.1 github.com/pkg/errors v0.9.1 + golang.org/x/sync v0.0.0-20210220032951-036812b2e83c google.golang.org/protobuf v1.25.0 ) diff --git a/go.sum b/go.sum index 17f7964d..d875a184 100644 --- a/go.sum +++ b/go.sum @@ -52,6 +52,8 @@ golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAG golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ= +golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=