zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: mapreduce generic version (#2827)

Browse Source 
* feat: mapreduce generic version

* fix: gateway mr type issue

---------

Co-authored-by: kevin.wan <kevin.wan@yijinin.com>
This commit is contained in:
Kevin Wan
2023-01-29 18:01:23 +08:00
committed by GitHub
parent 413ee919e6
commit 464ed51728
8 changed files with 166 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
core/mr/mapreduce.go
View File

@@ -7,7 +7,6 @@ import (
"sync/atomic"
"github.com/zeromicro/go-zero/core/errorx"
"github.com/zeromicro/go-zero/core/lang"
)
const (
@@ -24,30 +23,30 @@ var (
type (
// ForEachFunc is used to do element processing, but no output.
ForEachFunc func(item any)
ForEachFunc[T any] func(item T)
// GenerateFunc is used to let callers send elements into source.
GenerateFunc func(source chan<- any)
GenerateFunc[T any] func(source chan<- T)
// MapFunc is used to do element processing and write the output to writer.
MapFunc func(item any, writer Writer)
MapFunc[T, U any] func(item T, writer Writer[U])
// MapperFunc is used to do element processing and write the output to writer,
// use cancel func to cancel the processing.
MapperFunc func(item any, writer Writer, cancel func(error))
MapperFunc[T, U any] func(item T, writer Writer[U], cancel func(error))
// ReducerFunc is used to reduce all the mapping output and write to writer,
// use cancel func to cancel the processing.
ReducerFunc func(pipe <-chan any, writer Writer, cancel func(error))
ReducerFunc[U, V any] func(pipe <-chan U, writer Writer[V], cancel func(error))
// VoidReducerFunc is used to reduce all the mapping output, but no output.
// Use cancel func to cancel the processing.
VoidReducerFunc func(pipe <-chan any, cancel func(error))
VoidReducerFunc[U any] func(pipe <-chan U, cancel func(error))
// Option defines the method to customize the mapreduce.
Option func(opts *mapReduceOptions)
mapperContext struct {
mapperContext[T, U any] struct {
ctx context.Context
mapper MapFunc
source <-chan any
mapper MapFunc[T, U]
source <-chan T
panicChan *onceChan
collector chan<- any
doneChan <-chan lang.PlaceholderType
collector chan<- U
doneChan <-chan struct{}
workers int
}
@@ -57,8 +56,8 @@ type (
}
// Writer interface wraps Write method.
Writer interface {
Write(v any)
Writer[T any] interface {
Write(v T)
}
)
@@ -68,12 +67,11 @@ func Finish(fns ...func() error) error {
return nil
}
return MapReduceVoid(func(source chan<- any) {
return MapReduceVoid(func(source chan<- func() error) {
for _, fn := range fns {
source <- fn
}
}, func(item any, writer Writer, cancel func(error)) {
fn := item.(func() error)
}, func(fn func() error, writer Writer[any], cancel func(error)) {
if err := fn(); err != nil {
cancel(err)
}
@@ -87,27 +85,26 @@ func FinishVoid(fns ...func()) {
return
}
ForEach(func(source chan<- any) {
ForEach(func(source chan<- func()) {
for _, fn := range fns {
source <- fn
}
}, func(item any) {
fn := item.(func())
}, func(fn func()) {
fn()
}, WithWorkers(len(fns)))
}
// ForEach maps all elements from given generate but no output.
func ForEach(generate GenerateFunc, mapper ForEachFunc, opts ...Option) {
func ForEach[T any](generate GenerateFunc[T], mapper ForEachFunc[T], opts ...Option) {
options := buildOptions(opts...)
panicChan := &onceChan{channel: make(chan any)}
source := buildSource(generate, panicChan)
collector := make(chan any)
done := make(chan lang.PlaceholderType)
done := make(chan struct{})
go executeMappers(mapperContext{
go executeMappers(mapperContext[T, any]{
ctx: options.ctx,
mapper: func(item any, _ Writer) {
mapper: func(item T, _ Writer[any]) {
mapper(item)
},
source: source,
@@ -131,26 +128,26 @@ func ForEach(generate GenerateFunc, mapper ForEachFunc, opts ...Option) {
// MapReduce maps all elements generated from given generate func,
// and reduces the output elements with given reducer.
func MapReduce(generate GenerateFunc, mapper MapperFunc, reducer ReducerFunc,
opts ...Option) (any, error) {
func MapReduce[T, U, V any](generate GenerateFunc[T], mapper MapperFunc[T, U], reducer ReducerFunc[U, V],
opts ...Option) (V, error) {
panicChan := &onceChan{channel: make(chan any)}
source := buildSource(generate, panicChan)
return mapReduceWithPanicChan(source, panicChan, mapper, reducer, opts...)
}
// MapReduceChan maps all elements from source, and reduce the output elements with given reducer.
func MapReduceChan(source <-chan any, mapper MapperFunc, reducer ReducerFunc,
opts ...Option) (any, error) {
func MapReduceChan[T, U, V any](source <-chan T, mapper MapperFunc[T, U], reducer ReducerFunc[U, V],
opts ...Option) (V, error) {
panicChan := &onceChan{channel: make(chan any)}
return mapReduceWithPanicChan(source, panicChan, mapper, reducer, opts...)
}
// mapReduceWithPanicChan maps all elements from source, and reduce the output elements with given reducer.
func mapReduceWithPanicChan(source <-chan any, panicChan *onceChan, mapper MapperFunc,
reducer ReducerFunc, opts ...Option) (any, error) {
func mapReduceWithPanicChan[T, U, V any](source <-chan T, panicChan *onceChan, mapper MapperFunc[T, U],
reducer ReducerFunc[U, V], opts ...Option) (val V, err error) {
options := buildOptions(opts...)
// output is used to write the final result
output := make(chan any)
output := make(chan V)
defer func() {
// reducer can only write once, if more, panic
for range output {
@@ -159,12 +156,12 @@ func mapReduceWithPanicChan(source <-chan any, panicChan *onceChan, mapper Mappe
}()
// collector is used to collect data from mapper, and consume in reducer
collector := make(chan any, options.workers)
collector := make(chan U, options.workers)
// if done is closed, all mappers and reducer should stop processing
done := make(chan lang.PlaceholderType)
done := make(chan struct{})
writer := newGuardedWriter(options.ctx, output, done)
var closeOnce sync.Once
// use atomic.Value to avoid data race
// use atomic type to avoid data race
var retErr errorx.AtomicError
finish := func() {
closeOnce.Do(func() {
@@ -195,9 +192,9 @@ func mapReduceWithPanicChan(source <-chan any, panicChan *onceChan, mapper Mappe
reducer(collector, writer, cancel)
}()
go executeMappers(mapperContext{
go executeMappers(mapperContext[T, U]{
ctx: options.ctx,
mapper: func(item any, w Writer) {
mapper: func(item T, w Writer[U]) {
mapper(item, w, cancel)
},
source: source,
@@ -210,26 +207,29 @@ func mapReduceWithPanicChan(source <-chan any, panicChan *onceChan, mapper Mappe
select {
case <-options.ctx.Done():
cancel(context.DeadlineExceeded)
return nil, context.DeadlineExceeded
err = context.DeadlineExceeded
case v := <-panicChan.channel:
// drain output here, otherwise for loop panic in defer
drain(output)
panic(v)
case v, ok := <-output:
if err := retErr.Load(); err != nil {
return nil, err
if e := retErr.Load(); e != nil {
err = e
} else if ok {
return v, nil
val = v
} else {
return nil, ErrReduceNoOutput
err = ErrReduceNoOutput
}
}
return
}
// MapReduceVoid maps all elements generated from given generate,
// and reduce the output elements with given reducer.
func MapReduceVoid(generate GenerateFunc, mapper MapperFunc, reducer VoidReducerFunc, opts ...Option) error {
_, err := MapReduce(generate, mapper, func(input <-chan any, writer Writer, cancel func(error)) {
func MapReduceVoid[T, U any](generate GenerateFunc[T], mapper MapperFunc[T, U],
reducer VoidReducerFunc[U], opts ...Option) error {
_, err := MapReduce(generate, mapper, func(input <-chan U, writer Writer[any], cancel func(error)) {
reducer(input, cancel)
}, opts...)
if errors.Is(err, ErrReduceNoOutput) {
@@ -266,8 +266,8 @@ func buildOptions(opts ...Option) *mapReduceOptions {
return options
}
func buildSource(generate GenerateFunc, panicChan *onceChan) chan any {
source := make(chan any)
func buildSource[T any](generate GenerateFunc[T], panicChan *onceChan) chan T {
source := make(chan T)
go func() {
defer func() {
if r := recover(); r != nil {
@@ -283,13 +283,13 @@ func buildSource(generate GenerateFunc, panicChan *onceChan) chan any {
}
// drain drains the channel.
func drain(channel <-chan any) {
func drain[T any](channel <-chan T) {
// drain the channel
for range channel {
}
}
func executeMappers(mCtx mapperContext) {
func executeMappers[T, U any](mCtx mapperContext[T, U]) {
var wg sync.WaitGroup
defer func() {
wg.Wait()
@@ -298,7 +298,7 @@ func executeMappers(mCtx mapperContext) {
}()
var failed int32
pool := make(chan lang.PlaceholderType, mCtx.workers)
pool := make(chan struct{}, mCtx.workers)
writer := newGuardedWriter(mCtx.ctx, mCtx.collector, mCtx.doneChan)
for atomic.LoadInt32(&failed) == 0 {
select {
@@ -306,7 +306,7 @@ func executeMappers(mCtx mapperContext) {
return
case <-mCtx.doneChan:
return
case pool <- lang.Placeholder:
case pool <- struct{}{}:
item, ok := <-mCtx.source
if !ok {
<-pool
@@ -346,22 +346,21 @@ func once(fn func(error)) func(error) {
}
}
type guardedWriter struct {
type guardedWriter[T any] struct {
ctx context.Context
channel chan<- any
done <-chan lang.PlaceholderType
channel chan<- T
done <-chan struct{}
}
func newGuardedWriter(ctx context.Context, channel chan<- any,
done <-chan lang.PlaceholderType) guardedWriter {
return guardedWriter{
func newGuardedWriter[T any](ctx context.Context, channel chan<- T, done <-chan struct{}) guardedWriter[T] {
return guardedWriter[T]{
ctx: ctx,
channel: channel,
done: done,
}
}
func (gw guardedWriter) Write(v any) {
func (gw guardedWriter[T]) Write(v T) {
select {
case <-gw.ctx.Done():
return