getfriends_errgroup.go

func GetFriends(ctx context.Context, user int64) (map[string]*User, error) {
  g, ctx := errgroup.WithContext(ctx)
  friendIds := make(chan int64)

  // Produce
  g.Go(func() error {
     defer close(friendIds)
     for it := GetFriendIds(user); ; {
        if id, err := it.Next(ctx); err != nil {
           if err == io.EOF {
              return nil
           }
           return fmt.Errorf("GetFriendIds %d: %s", user, err)
        } else {
           friendIds <- id
        }
     }
  })

  friends := make(chan *User)

  // Map
  workers := int32(nWorkers)
  for i := 0; i < nWorkers; i++ {
     g.Go(func() error {
        defer func() {
           // Last one out closes shop
           if atomic.AddInt32(&workers, -1) == 0 {
              close(friends)
           }
        }()

        for id := range friendIds {
           if friend, err := GetUserProfile(ctx, id); err != nil {
              return fmt.Errorf("GetUserProfile %d: %s", user, err)
           } else {
              friends <- friend
           }
        }
        return nil
     })
  }

  // Reduce
  ret := map[string]*User{}
  g.Go(func() error {
     for friend := range friends {
        ret[friend.Name] = friend
     }
     return nil
  })

  // Return the final result, and the error result from the subtask group.
  return ret, g.Wait()
}