feat: go goroutine wrapper

go-god · Jan 9, 2022 · d41dd5e · d41dd5e
1 parent d970a56
commit d41dd5e
Show file tree

Hide file tree

Showing 11 changed files with 334 additions and 0 deletions.
diff --git a/chanwrap/chan_wrap_impl.go b/chanwrap/chan_wrap_impl.go
@@ -0,0 +1,59 @@
+package chanwrap
+
+import (
+	"github.com/go-god/wrapper"
+	"github.com/go-god/wrapper/grecover"
+)
+
+var est = struct{}{}
+
+// WrapImpl wrapper impl
+type WrapImpl struct {
+	bufNum       int
+	bufCh        chan struct{}
+	recoveryFunc func()
+}
+
+// New create wrapImpl entity
+func New(c int) wrapper.Wrapper {
+	w := &WrapImpl{
+		bufNum:       c,
+		bufCh:        make(chan struct{}, c),
+		recoveryFunc: grecover.DefaultRecovery,
+	}
+
+	return w
+}
+
+// Wrap exec func in goroutine without recover catch
+func (c *WrapImpl) Wrap(fn func()) {
+	go func() {
+		defer c.done()
+		fn()
+	}()
+}
+
+// WrapWithRecover safely execute func in goroutine
+func (c *WrapImpl) WrapWithRecover(fn func()) {
+	go func() {
+		defer c.recoveryFunc()
+		defer c.done()
+		fn()
+	}()
+}
+
+// Wait wait all goroutine finish
+func (c *WrapImpl) Wait() {
+	for i := 0; i < c.bufNum; i++ {
+		<-c.bufCh
+	}
+}
+
+// WithRecover set recover func
+func (c *WrapImpl) WithRecover(recoveryFunc func()) {
+	c.recoveryFunc = recoveryFunc
+}
+
+func (c *WrapImpl) done() {
+	c.bufCh <- est
+}
diff --git a/chanwrap/chan_wrap_test.go b/chanwrap/chan_wrap_test.go
@@ -0,0 +1,20 @@
+package chanwrap
+
+import (
+	"log"
+	"testing"
+)
+
+func TestWrapper(t *testing.T) {
+	var wg = New(2)
+	wg.Wrap(func() {
+		log.Println("1111")
+	})
+	wg.WrapWithRecover(func() {
+		log.Println(2222)
+		panic("mock panic test")
+	})
+	wg.Wait()
+
+	// time.Sleep(10 * time.Second)
+}
diff --git a/example/main.go b/example/main.go
@@ -0,0 +1,43 @@
+package main
+
+import (
+	"log"
+
+	"github.com/go-god/wrapper/factory"
+)
+
+func main() {
+	wrapper := factory.New(factory.ChWrapper, 2)
+	wrapper.Wrap(func() {
+		log.Println("chan wrapper: 1111")
+	})
+	wrapper.Wrap(func() {
+		log.Println("chan wrapper: 2222")
+	})
+
+	wrapper.Wait()
+
+	// factory.WgWrapper No need to pass second parameter.
+	wg := factory.New(factory.WgWrapper)
+	wg.Wrap(func() {
+		log.Println("wg wrapper:1111")
+	})
+	wg.Wrap(func() {
+		log.Println("wg wrapper:2222")
+	})
+
+	wg.WrapWithRecover(func() {
+		log.Println("wg wrapper:3333")
+		panic("mock panic:abc")
+	})
+
+	wg.Wait()
+}
+
+// output:
+// 2022/01/09 22:47:48 chan wrapper: 2222
+// 2022/01/09 22:47:48 chan wrapper: 1111
+// 2022/01/09 22:47:48 wg wrapper:3333
+// 2022/01/09 22:47:48 wg wrapper:1111
+// 2022/01/09 22:47:48 wrapper exec recover:mock panic:abc
+// 2022/01/09 22:47:48 wg wrapper:2222
diff --git a/factory/wrap_factory.go b/factory/wrap_factory.go
@@ -0,0 +1,44 @@
+package factory
+
+import (
+	"github.com/go-god/wrapper"
+	"github.com/go-god/wrapper/chanwrap"
+	"github.com/go-god/wrapper/waitgroup"
+)
+
+type constructor func() wrapper.Wrapper
+
+const (
+	// WgWrapper waitGroup wrapper
+	WgWrapper = "wg"
+	// ChWrapper chan wrapper
+	ChWrapper = "ch"
+)
+
+var wrapperMap = map[string]constructor{
+	WgWrapper: waitgroup.New,
+	ChWrapper: nil,
+}
+
+// New create wrapper interface
+func New(name string, c ...int) wrapper.Wrapper {
+	if name == ChWrapper && len(c) > 0 && c[0] > 0 {
+		return chanwrap.New(c[0])
+	}
+
+	if w, ok := wrapperMap[name]; ok {
+		return w()
+	}
+
+	panic("wrapper type not exists")
+}
+
+// Register register gdi.Injector
+func Register(name string, c constructor) {
+	_, ok := wrapperMap[name]
+	if ok {
+		panic("registered injector already exists")
+	}
+
+	wrapperMap[name] = c
+}
diff --git a/factory/wrap_factory_test.go b/factory/wrap_factory_test.go
@@ -0,0 +1,15 @@
+package factory
+
+import (
+	"log"
+	"testing"
+)
+
+func TestNew(t *testing.T) {
+	wrap := New("wg", 1)
+	wrap.Wrap(func() {
+		log.Println(1111)
+	})
+
+	wrap.Wait()
+}
diff --git a/go.mod b/go.mod
@@ -0,0 +1,3 @@
+module github.com/go-god/wrapper
+
+go 1.16
diff --git a/grecover/recovery.go b/grecover/recovery.go
@@ -0,0 +1,10 @@
+package grecover
+
+import "log"
+
+// DefaultRecovery default recover func.
+func DefaultRecovery() {
+	if e := recover(); e != nil {
+		log.Printf("wrapper exec recover:%v", e)
+	}
+}
diff --git a/readme.md b/readme.md
@@ -0,0 +1,54 @@
+# goroutine wrapper
+
+    The goroutine execution wrapper
+    the user only needs to pass the corresponding func to 
+    the wrap method to execute (supports safe execution), 
+    and supports waiting for all coroutines to complete execution.
+
+# Instructions for use
+
+```go
+package main
+
+import (
+	"log"
+
+	"github.com/go-god/wrapper/factory"
+)
+
+func main() {
+	wrapper := factory.New(factory.ChWrapper, 2)
+	wrapper.Wrap(func() {
+		log.Println("chan wrapper: 1111")
+	})
+	wrapper.Wrap(func() {
+		log.Println("chan wrapper: 2222")
+	})
+
+	wrapper.Wait()
+
+	// factory.WgWrapper No need to pass second parameter.
+	wg := factory.New(factory.WgWrapper)
+	wg.Wrap(func() {
+		log.Println("wg wrapper:1111")
+	})
+	wg.Wrap(func() {
+		log.Println("wg wrapper:2222")
+	})
+
+	wg.WrapWithRecover(func() {
+		log.Println("wg wrapper:3333")
+		panic("mock panic:abc")
+	})
+
+	wg.Wait()
+}
+
+// output:
+// 2022/01/09 22:47:48 chan wrapper: 2222
+// 2022/01/09 22:47:48 chan wrapper: 1111
+// 2022/01/09 22:47:48 wg wrapper:3333
+// 2022/01/09 22:47:48 wg wrapper:1111
+// 2022/01/09 22:47:48 wrapper exec recover:mock panic:abc
+// 2022/01/09 22:47:48 wg wrapper:2222
+```
diff --git a/waitgroup/waitgroup_impl.go b/waitgroup/waitgroup_impl.go
@@ -0,0 +1,46 @@
+package waitgroup
+
+import (
+	"sync"
+
+	"github.com/go-god/wrapper"
+	"github.com/go-god/wrapper/grecover"
+)
+
+// WrapImpl sync.WaitGroup wrap impl
+type WrapImpl struct {
+	sync.WaitGroup
+	recoveryFunc func()
+}
+
+// New create wrapper entity
+func New() wrapper.Wrapper {
+	w := &WrapImpl{
+		recoveryFunc: grecover.DefaultRecovery,
+	}
+
+	return w
+}
+
+// Wrap fn func in goroutine to run
+func (w *WrapImpl) Wrap(fn func()) {
+	w.Add(1)
+	go func() {
+		defer w.Done()
+		fn()
+	}()
+}
+
+// WrapWithRecover exec func with recover
+func (w *WrapImpl) WrapWithRecover(fn func()) {
+	w.Add(1)
+	go func() {
+		defer w.recoveryFunc()
+		defer w.Done()
+		fn()
+	}()
+}
+
+func (w *WrapImpl) WithRecover(recoveryFunc func()) {
+	w.recoveryFunc = recoveryFunc
+}
diff --git a/waitgroup/waitgroup_test.go b/waitgroup/waitgroup_test.go
@@ -0,0 +1,24 @@
+package waitgroup
+
+import (
+	"log"
+	"testing"
+)
+
+func mockRecovery() {
+	if err := recover(); err != nil {
+		log.Printf("exec recover:%v\n", err)
+	}
+}
+
+func TestWrapper(t *testing.T) {
+	var wg = New(WithRecovery(mockRecovery))
+	wg.Wrap(func() {
+		log.Println("1111")
+	})
+	wg.WrapWithRecover(func() {
+		log.Println(2222)
+		panic("mock panic test")
+	})
+	wg.Wait()
+}
diff --git a/wrapper.go b/wrapper.go
@@ -0,0 +1,16 @@
+package wrapper
+
+// Wrapper wrap goroutine to run
+type Wrapper interface {
+	// Wrap exec func in goroutine without recover catch
+	Wrap(fn func())
+
+	// WrapWithRecover safely execute func in goroutine
+	WrapWithRecover(fn func())
+
+	// WithRecover set recover func
+	WithRecover(recoveryFunc func())
+
+	// Wait wait all goroutine finish
+	Wait()
+}