golang手动实现error group

我就是我 2022-11-03 04:15 143阅读 0赞

手动实现了下go error group，并发请求、请求数量限制、出错时cancel掉所有task，以下为示例，要求task是可cancel的！

package main
    
    import (
    	"context"
    	"errors"
    	"fmt"
    	"sync"
    	"sync/atomic"
    	"time"
    )
    
    const (
    	M = 2
    	N = 8
    )
    
    func main() {
    	ctx, cancel := context.WithTimeout(context.Background(), time.Second*50)
    	defer cancel()
    
    	result := make([]int, N+1)
    	errCh := make(chan error, 1)
    	var firstSendErr int32
    	wg := new(sync.WaitGroup)
    	done := make(chan struct{}, 1)
    	limit := make(chan struct{}, M)
    
    	for i := 1; i <= N; i++ {
    		limit <- struct{}{}
    		var quit bool
    		select {
    		case <-ctx.Done(): // context已经被cancel，不需要起新的goroutine了
    			quit = true
    		default:
    		}
    		if quit {
    			break
    		}
    		wg.Add(1)
    		go func(x int) {
    			defer func() {
    				wg.Done()
    				<-limit
    			}()
    			if ret, err := doTask(ctx, x); err != nil {
    				if atomic.CompareAndSwapInt32(&firstSendErr, 0, 1) {
    					errCh <- err
    					cancel() // cancel其他的请求
    				}
    			} else {
    				result[x] = ret
    			}
    		}(i)
    	}
    
    	go func() {
    		wg.Wait()
    		close(done)
    	}()
    
    	select {
    	case err := <-errCh:
    		handleErr(err, result[1:])
    		<-done
    	case <-done:
    		if len(errCh) > 0 {
    			err := <-errCh
    			handleErr(err, result[1:])
    			return
    		}
    		fmt.Println("success handle all task:", result[1:])
    	}
    }
    
    func handleErr(err error, result []int) {
    	fmt.Println("task err occurs: ", err, "result", result)
    }
    
    func doTask(ctx context.Context, i int) (ret int, err error) {
    	fmt.Println("task start", i)
    	defer func() {
    		fmt.Println("task done", i, "err", err)
    	}()
    	select {
    	case <-time.After(time.Second * time.Duration(i)): // 模拟处理任务时间
    	case <-ctx.Done(): // 处理任务要支持被context cancel，不然就一直等到处理完再返回了。
    		fmt.Println("task canceled", i)
    		return -1, ctx.Err()
    	}
    	if i == 6 { // 模拟出现错误
    		return -1, errors.New("err test")
    	}
    
    	return i, nil
    }