Asynchronous Go iterator

async_iter.go

// This example demonstrates getting a series of results and an error from a
// goroutine, with the ability to cancel the goroutine. The goroutine can be
// canceled by stopping iteration (break out of the iteration loop), or by
// canceling a context, as may be the case if a context from outside of the
// iteration logic is provided.
//
// The first part of the example shows using channels to read the results and
// error, and using a context to cancel the goroutine before all results are
// received.
//
// The second part shows how this can be done using an iterator to return
// results and an error, and how the goroutine is canceled if iteration is
// stopped before all results are read.
//
// While both approaches do the same work, the iterator provides a cleaner
// interface that is easier to use and harder to misuse.
//
// https://go.dev/play/p/8wjN1p1h9j5
package main

import (
	"context"
	"fmt"
	"iter"
)

const cancelAt = "wed"
const errorAt = "fri"

func main() {
	// Note for both examples: there is no need to check if the context
	// has been canceled, outside of asyncTask, because asyncTask checks
	// and will return ctx.Err() if it has.
	
	// Raw channels
	// ============

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	results, asyncErr := asyncTask(ctx)
	for result := range results {
		fmt.Println("Day:", result)
		if result == cancelAt {
			cancel()       // cancel asyncTask1
			asyncErr = nil // canceled, so do not get error
			fmt.Println(result, "is may last day")
			break
		}
	}
	if asyncErr != nil {
		if err := <-asyncErr; err != nil {
			fmt.Println("Error1:", err)
		}
	}

	fmt.Println("----------------------------------------")

	// Iterator
	// ========

	for result, err := range asyncIter(context.Background()) {
		if err != nil {
			fmt.Println("Error2:", err)
			break
		}
		fmt.Println("Day:", result)
		if result == cancelAt {
			fmt.Println(result, "is may last day")
			break
		}
	}
}

// asyncTask returns days of the week over a channel from a goroutine. It sends
// an error when the day is equal to errorAt.
func asyncTask(ctx context.Context) (<-chan string, <-chan error) {
	results := make(chan string)
	asyncErr := make(chan error, 1)

	go func() {
		defer close(results)
		defer close(asyncErr)

		days := []string{"mon", "tue", "wed", "thu", "fri", "sat", "sun"}
		for _, day := range days {
			if day == errorAt {
				asyncErr <- fmt.Errorf("there is a problem on %s", day)
				return
			}
			select {
			case results <- day:
			case <-ctx.Done():
				asyncErr <- ctx.Err()
				return
			}
		}
	}()

	return results, asyncErr
}

// asyncIter wraps asyncTask in an iterator. If iteration in canceled, then the
// task is canceled.
func asyncIter(ctx context.Context) iter.Seq2[string, error] {
	return func(yield func(string, error) bool) {
		ctx, cancel := context.WithCancel(ctx)
		defer cancel() // cancel task when done iterating
		results, asyncErr := asyncTask(ctx)

		for result := range results {
			if !yield(result, nil) {
				return
			}
		}
		if err := <-asyncErr; err != nil {
			yield("", err)
		}
	}
}