Go Generator with Context

main.go

package main

import (
	"context"
	"fmt"
	"math"
	"time"
)

func main() {
	// ch := make(chan int)
	ctx := context.Background()
	timeout := 5 * time.Minute
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	start := time.Now()
	ch := PrimeGeneratorHidden(ctx)

	var p int
	var n int

	for p = range ch {
		n++
		if n%1000000 == 0 {
			fmt.Printf("%d: %d\n", n, p)
		}
	}

	elapsed := time.Since(start)
	fmt.Printf("Found prime #%d: %d, after %s.\n", n, p, elapsed)
}

func PrimeGeneratorHidden(ctx context.Context) <-chan int {
	ch := make(chan int)

	go PrimeGenerator(ctx, ch)

	return ch
}

func PrimeGenerator(ctx context.Context, ch chan<- int) {
	primes := make([]int, 0)
	ch <- 2
	for i := 3; ; i += 2 {
		isprime := true
		iSqrt := int(math.Floor(math.Sqrt(float64(i))))
		for _, p := range primes {
			if i%p == 0 {
				isprime = false
				break
			}
			if p > iSqrt {
				break
			}
		}
		if isprime {
			select {
			case ch <- i:
				primes = append(primes, i)
			case <-ctx.Done():
				close(ch)
				return
			}
		}
	}
}