等待两个结果同时超时

Question

Use case

I'd like to run two queries against a database in parallel and return after a maximum time of 600ms whatever I have fetched to that point. I am struggling with implementing the concurrency for this requirement.

Code

func (s *Service) GetCustomerStats(ctx context.Context, customerUUID string) *CustomerStats {
    stats := &CustomerStats{
        CustomerUUID: customerUUID,
        Type:         "ERROR",
        OrderCount:   "ERROR",
    }

    var wg sync.WaitGroup
    var mu sync.Mutex

    // Get order count
    wg.Add(1)
    go func() {
        defer wg.Done()
        orderCount, err := s.Storage.GetOrderCount(ctx, customerUUID)
        if err != nil {
            return
        }

        mu.Lock()
        stats.OrderCount = strconv.Itoa(orderCount)
        if orderCount == 0 {
            stats.OrderCount = "NA"
        }
        mu.Unlock()
    }()

    // Get customer type
    wg.Add(1)
    go func() {
        defer wg.Done()
        type, err := s.Storage.GetCustomerType(ctx, customerUUID)
        if err != nil {
            return
        }

        mu.Lock()
        stats.Type = strconv.Itoa(type)
        mu.Unlock()
    }()

    wg.Wait()

    return stats
}

The problem

The context I pass into that function has a timeout of 600ms defined. I pass it on to the storage repo and the Database driver uses it as well, but it does not guarantee it will respond within that time as it does schedule some retries under the hood.

However I must ensure that this function returns within the passed context timeout (600ms). I am currently using a waitgroup to await the results but I wouldn't know how to return stats once the context is done.

Basically I am looking for something like this. My research indicates that I should probably use channels which signal that the work is done but I am not sure how I would implement that so that it's simple code.

    select {
    case wg.Wait()
        return stats
    case <-ctx.Done()
        return stats
    }

Answer 1

The way you plan to select on ctx.Done() looks correct.
It's the way you work with your mutable state that is wrong, in my opinion.

Try something like this:

 var state = State{}
 select {
    case type <- typeChan
        stats.Type = type
        if (stats.OrderCount != nil) {
           return stats
        }
    case count <- countChan
        stats.OrderCount = count
        if (stats.Type != nil) {
           return stats
        }
    case <-ctx.Done()
        return stats
    }

Now your functions should look like this:

go func() {        
    orderCount, err := s.Storage.GetOrderCount(ctx, customerUUID)
    if err != nil {
        return // Here you probably want to have errChan
    }

    if orderCount == 0 {
        countChan <- "NA"
    } else {
        countChan <- strconv.Itoa(orderCount)
    }        
}()

This is all a bit sketchy, since your example is quite complex, but should give you the direction to follow.