golang channel 怎么判断是否读取完成(2)？

Question

需求是判断p是否满足p=a^2+b^2+c^2

package main

import (
    "fmt"
    "math"
)

func abcasync(p int) bool {
    na := int(math.Sqrt(float64(p))) + 2
    channel := make(chan bool)
    for a := 0; a < na; a++ {
        for b := 0; b < na; b++ {
            go func() {
                for c := 0; c < na; c++ {
                    if a*a+b*b+c*c == p {
                        channel <- true
                    }
                }
            }()

        }
    }
    return <-channel
}

func main() {
    // d := abcasync(41)
    d := abcasync(23)
    fmt.Println(d)
}

这里要求，一旦满足if条件，就退出程序，以避免浪费计算时间，否则就要遍历所有的情况。那么，如何判断所有的协程都执行完了呢？
会报错，请教如何解决？

Answer 1

package main

import (
    "fmt"
    "math"
    "sync"
    "sync/atomic"
)

func abcasync(p int) bool {
    na := int(math.Sqrt(float64(p))) + 2
    var waiteRet sync.WaitGroup
    var find int32 = 0
loop:
    for a := 0; a < na; a++ {
        for b := 0; b < na; b++ {
            if find == 1 {
                break loop
            }
            waiteRet.Add(1)
            go func(a, b int) {
                defer waiteRet.Done()
                for c := 0; c < na; c++ {
                    if find == 1 {
                        break
                    }
                    if a*a+b*b+c*c == p {
                        if atomic.CompareAndSwapInt32(&find, 0, 1) {
                            fmt.Println(a, b, c)
                            break
                        }
                    }
                }
            }(a, b)

        }
    }
    waiteRet.Wait()
    return atomic.LoadInt32(&find) == 1
}

应该能满足你的需求，中间beak判断没有使用atomic。使用chan，找不到的情况下会死锁啊哥哥...还有很多重复计算...

Answer 2

package main

import (
    "fmt"
    "math"
    "runtime"
    "sync"
)

type S struct {
    a int
    b int
    c int
    p int
}

func NewS(pa, pb, pc, pp int) *S {
    return &S{
        a: pa,
        b: pb,
        c: pc,
        p: pp,
    }
}

func (s *S) isMatch() bool {
    return s.a*s.a+s.b*s.b+s.c*s.c == s.p
}

func abcasync(ch chan *S, chResult chan *S, p int) *S {
    na := int(math.Sqrt(float64(p))) + 2
    for a := 0; a < na; a++ {
        for b := 0; b < na; b++ {
            for c := 0; c < na; c++ {
                select {
                case r := <-chResult:
                    close(ch)
                    return r
                default:
                    ch <- (NewS(a, b, c, p))
                }
            }
        }
    }
    return nil
}

func main() {
    pNum := runtime.NumCPU()
    ch := make(chan *S, 100)
    chResult := make(chan *S)

    var wg sync.WaitGroup
    wg.Add(pNum)
    for i := 0; i < pNum; i++ {
        go func() {
            defer wg.Done()
            for s := range ch {
                if s.isMatch() {
                    chResult <- s
                }
            }
        }()
    }
    d := abcasync(ch, chResult, 10000)
    if d == nil {
        fmt.Println("not found")
    } else {
        fmt.Println(d)
    }
    wg.Wait()
}

计算密集型的，携程多了不管用，你可以增加协程数benchmark一下