在并发Go应用中保持顺序【译】

// Map transforms items from the input channel using n goroutines, and the
// provided function f. Returns a new channel with transformed items.
func Map[A, B any](in <-chan A, n int, f func(A) B) <-chan B {
	out := make(chan B)
	Loop(in, n, out, func(a A) {
		out <- f(a)
	})
	return out
}

// Loop is a worker pool implementation. It calls function f for each 
// item from the input channel using n goroutines. This is a non-blocking function 
// that signals completion by closing the done channel when all work is finished.
func Loop[A, B any](in <-chan A, n int, done chan<- B, f func(A)) {
	var wg sync.WaitGroup

	for i := 0; i < n; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for a := range in {
				f(a)
			}
		}()
	}

	go func() {
		wg.Wait()
		if done != nil {
			close(done)
		}
	}()
}

// Discard is a non-blocking function that consumes and discards
// all items from the input channel
func Discard[A any](in <-chan A) {
	go func() {
		for range in {
			// Discard the value
		}
	}()
}

func BenchmarkMap(b *testing.B) {
	for _, n := range []int{1, 2, 4, 8, 12, 50} {
		b.Run(fmt.Sprint("n=", n), func(b *testing.B) {
			in := make(chan int)
			defer close(in)
			out := Map(in, n, func(a int) int {
				//time.Sleep(50 * time.Microsecond)
				return a // no-op: just return the original value
			})
			Discard(out)

			b.ReportAllocs()
			b.ResetTimer()

			for i := 0; i < b.N; i++ {
				in <- 10 // write something to the in chan
			}
		})
	}
}

func OrderedMap1[A, B any](in <-chan A, n int, f func(A) B) <-chan B {
	type Job struct {
		Item    A
		ReplyTo chan B
	}

	// Packer goroutine.
	// `jobs` chan will be processed by the pool
	// `replies` chan will be consumed by unpacker goroutine
	jobs := make(chan Job)
	replies := make(chan chan B, n)
	go func() {
		for item := range in {
			replyTo := make(chan B)
			jobs <- Job{Item: item, ReplyTo: replyTo}
			replies <- replyTo
		}
		close(jobs)
		close(replies)
	}()

	// Worker pool of n goroutines.
	// Sends results back via replyTo channels
	Loop[Job, any](jobs, n, nil, func(job Job) {
		job.ReplyTo <- f(job.Item) // Calculate the result and send it back
		close(job.ReplyTo)
	})

	// Unpacker goroutine.
	// Unpacks replyTo channels in order and sends results to the `out` channel
	out := make(chan B)
	go func() {
		defer close(out)
		for replyTo := range replies {
			result := <-replyTo
			out <- result
		}
	}()
	return out
}

func OrderedMap2[A, B any](in <-chan A, n int, f func(A) B) <-chan B {
	type Job struct {
		Item  A
		Index int
	}

	// Indexer goroutine.
	// Assign an index to each item from the input channel
	jobs := make(chan Job)
	go func() {
		i := 0
		for item := range in {
			jobs <- Job{Item: item, Index: i}
			i++
		}
		close(jobs)
	}()

	// Shared state.
	// Index of the next result that must be written to the output channel.
	nextIndex := 0
	cond := sync.NewCond(new(sync.Mutex))

	// Worker pool of n goroutines.
	out := make(chan B)
	Loop(jobs, n, out, func(job Job) {
		result := f(job.Item) // Calculate the result

		// Cond must be used with a locked mutex (see stdlib docs)
		cond.L.Lock()

		// wait until it's our turn to write the result
		for job.Index != nextIndex {
			cond.Wait()
		}

		// Write the result
		out <- result

		// Increment the index and notify all other workers
		nextIndex++
		cond.Broadcast()

		cond.L.Unlock()
	})

	return out
}

func OrderedMap3[A, B any](in <-chan A, n int, f func(A) B) <-chan B {
	type Job[A any] struct {
		Item         A
		CanWrite     chan struct{}
		NextCanWrite chan struct{} // canWrite channel of the next job
	}

	// Linker goroutine:
	// Builds a chain of jobs where each has a CanWrite channel attached.
	// Additionally, each job knows about the CanWrite channel of the next job in the chain.
	jobs := make(chan Job[A])
	go func() {
		defer close(jobs)

		var canWrite, nextCanWrite chan struct{}
		nextCanWrite = make(chan struct{}, 1)
		close(nextCanWrite) // the first job can write immediately

		for item := range in {
			canWrite, nextCanWrite = nextCanWrite, make(chan struct{}, 1)
			jobs <- Job[A]{item, canWrite, nextCanWrite}
		}
	}()

	// Worker pool of n goroutines.
	// Jobs pass the write permission along the chain.
	out := make(chan B)
	Loop(jobs, n, out, func(job Job[A]) {
		result := f(job.Item) // Calculate the result

		<-job.CanWrite          // Wait for the write permission
		out <- result           // Write to the output channel
		close(job.NextCanWrite) // Pass the permission to the next job
	})

	return out
}

// Package-level pool for canWrite channels
type chainedItem[A any] struct {
	Value        A
	CanWrite     chan struct{}
	NextCanWrite chan struct{} // canWrite channel for the next item
}

var canWritePool sync.Pool

func makeCanWriteChan() chan struct{} {
	ch := canWritePool.Get()
	if ch == nil {
		return make(chan struct{}, 1)
	}
	return ch.(chan struct{})
}

func releaseCanWriteChan(ch chan struct{}) {
	canWritePool.Put(ch)
}

func OrderedMap3a[A, B any](in <-chan A, n int, f func(A) B) <-chan B {
	type Job[A any] struct {
		Item         A
		CanWrite     chan struct{}
		NextCanWrite chan struct{} // canWrite channel of the next job
	}

	// Linker goroutine:
	// Builds a chain of jobs where each has a CanWrite channel attached.
	// Additionally, each job knows about the CanWrite channel of the next job in the chain.
	jobs := make(chan Job[A])
	go func() {
		defer close(jobs)

		var canWrite, nextCanWrite chan struct{}
		nextCanWrite = makeCanWriteChan()
		nextCanWrite <- struct{}{} // the first job can write immediately

		for item := range in {
			canWrite, nextCanWrite = nextCanWrite, makeCanWriteChan()
			jobs <- Job[A]{item, canWrite, nextCanWrite}
		}
	}()

	// Worker pool of n goroutines.
	// Jobs pass the write permission along the chain.
	out := make(chan B)
	Loop(jobs, n, out, func(job Job[A]) {
		result := f(job.Item) // Calculate the result

		<-job.CanWrite                    // Wait for the write permission
		out <- result                     // Write to the output channel
		releaseCanWriteChan(job.CanWrite) // Release our canWrite channel to the pool
		job.NextCanWrite <- struct{}{}    // Pass the permission to the next job
	})

	return out
}

func OrderedLoop[A, B any](in <-chan A, done chan<- B, n int, f func(a A, canWrite <-chan struct{})) {
	type Job[A any] struct {
		Item         A
		CanWrite     chan struct{}
		NextCanWrite chan struct{} // canWrite channel of the next job
	}

	// Linker goroutine:
	// Builds a chain of jobs where each has a CanWrite channel attached.
	// Additionally, each job knows about the CanWrite channel of the next job in the chain.
	jobs := make(chan Job[A])
	go func() {
		defer close(jobs)

		var canWrite, nextCanWrite chan struct{}
		nextCanWrite = makeCanWriteChan()
		nextCanWrite <- struct{}{} // the first job can write immediately

		for item := range in {
			canWrite, nextCanWrite = nextCanWrite, makeCanWriteChan()
			jobs <- Job[A]{item, canWrite, nextCanWrite}
		}
	}()

	// Worker pool of n goroutines.
	// Jobs pass the write permission along the chain.
	Loop(jobs, n, done, func(job Job[A]) {
		f(job.Item, job.CanWrite) // Do the work

		releaseCanWriteChan(job.CanWrite) // Release item's canWrite channel to the pool
		job.NextCanWrite <- struct{}{}    // Pass the permission to the next job
	})
}

OrderedLoop(in, out, n, func(a A, canWrite <-chan struct{}) {
	// [Do processing here]
	
	// Everything above this line is executed concurrently,
	// everything below it is executed sequentially and in order
	<-canWrite
	
	// [Write results somewhere]
})

func OrderedMap3b[A, B any](in <-chan A, n int, f func(A) B) <-chan B {
	out := make(chan B)
	OrderedLoop(in, out, n, func(a A, canWrite <-chan struct{}) {
		result := f(a)
		<-canWrite
		out <- result
	})
	return out
}

func OrderedSplit[A any](in <-chan A, n int, predicate func(A) bool) (<-chan A, <-chan A) {
	outTrue := make(chan A)
	outFalse := make(chan A)
	done := make(chan struct{})
	
	OrderedLoop(in, done, n, func(a A, canWrite <-chan struct{}) {
		shouldGoToTrue := predicate(a)
		<-canWrite
		if shouldGoToTrue {
			outTrue <- a
		} else {
			outFalse <- a
		}
	})
	
	go func() {
		<-done
		close(outTrue)
		close(outFalse)
	}()
	
	return outTrue, outFalse
}