feat: wrapper impl update

go-god · Apr 30, 2022 · bcab281 · bcab281
1 parent 3a9c686
commit bcab281
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 27 deletions.
diff --git a/chanwrap/chan_wrap_impl.go b/chanwrap/chan_wrap_impl.go
@@ -7,8 +7,10 @@ import (
 
 var est = struct{}{}
 
-// WrapImpl wrapper impl
-type WrapImpl struct {
+var _ wrapper.Wrapper = (*wrapImpl)(nil)
+
+// wrapImpl wrapper impl
+type wrapImpl struct {
 	bufCap       int
 	bufCh        chan struct{}
 	recoveryFunc func()
@@ -20,9 +22,8 @@ type WrapImpl struct {
 // Otherwise, after the Wait method is executed, some goroutines
 // will exit without execution.
 func New(opts ...wrapper.Option) wrapper.Wrapper {
-	w := &WrapImpl{}
-
-	var option = &wrapper.Options{}
+	w := &wrapImpl{}
+	option := &wrapper.Options{}
 	for _, o := range opts {
 		o(option)
 	}
@@ -39,15 +40,15 @@ func New(opts ...wrapper.Option) wrapper.Wrapper {
 }
 
 // Wrap exec func in goroutine without recover catch
-func (c *WrapImpl) Wrap(fn func()) {
+func (c *wrapImpl) Wrap(fn func()) {
 	go func() {
 		defer c.done()
 		fn()
 	}()
 }
 
 // WrapWithRecover safely execute func in goroutine
-func (c *WrapImpl) WrapWithRecover(fn func()) {
+func (c *wrapImpl) WrapWithRecover(fn func()) {
 	go func() {
 		defer c.recoveryFunc()
 		defer c.done()
@@ -56,12 +57,12 @@ func (c *WrapImpl) WrapWithRecover(fn func()) {
 }
 
 // Wait wait all goroutine finish
-func (c *WrapImpl) Wait() {
+func (c *wrapImpl) Wait() {
 	for i := 0; i < c.bufCap; i++ {
 		<-c.bufCh
 	}
 }
 
-func (c *WrapImpl) done() {
+func (c *wrapImpl) done() {
 	c.bufCh <- est
 }
diff --git a/factory/wrap_factory.go b/factory/wrap_factory.go
@@ -8,33 +8,36 @@ import (
 
 type constructor func(opts ...wrapper.Option) wrapper.Wrapper
 
+// WrapType wrap type
+type WrapType int
+
 const (
 	// WgWrapper waitGroup wrapper
-	WgWrapper = "wg"
+	WgWrapper WrapType = iota
 	// ChWrapper chan wrapper
-	ChWrapper = "ch"
+	ChWrapper
 )
 
-var wrapperMap = map[string]constructor{
+var wrapperMap = map[WrapType]constructor{
 	WgWrapper: waitgroup.New,
 	ChWrapper: chanwrap.New,
 }
 
 // New create wrapper interface
-func New(name string, opts ...wrapper.Option) wrapper.Wrapper {
-	if w, ok := wrapperMap[name]; ok {
+func New(wrapType WrapType, opts ...wrapper.Option) wrapper.Wrapper {
+	if w, ok := wrapperMap[wrapType]; ok {
 		return w(opts...)
 	}
 
 	panic("wrapper type not exists")
 }
 
 // Register register wrapper
-func Register(name string, c constructor) {
-	_, ok := wrapperMap[name]
+func Register(wrapType WrapType, c constructor) {
+	_, ok := wrapperMap[wrapType]
 	if ok {
 		panic("registered injector already exists")
 	}
 
-	wrapperMap[name] = c
+	wrapperMap[wrapType] = c
 }
diff --git a/factory/wrap_factory_test.go b/factory/wrap_factory_test.go
@@ -6,7 +6,7 @@ import (
 )
 
 func TestNew(t *testing.T) {
-	wrap := New("wg")
+	wrap := New(WgWrapper)
 	wrap.Wrap(func() {
 		log.Println(1111)
 	})

diff --git a/waitgroup/waitgroup_impl.go b/waitgroup/waitgroup_impl.go
@@ -7,19 +7,18 @@ import (
 	"github.com/go-god/wrapper/grecover"
 )
 
-// WrapImpl sync.WaitGroup wrap impl
-type WrapImpl struct {
+var _ wrapper.Wrapper = (*wrapImpl)(nil)
+
+// wrapImpl sync.WaitGroup wrap impl
+type wrapImpl struct {
 	sync.WaitGroup
 	recoveryFunc func()
 }
 
 // New create wrapper entity
 func New(opts ...wrapper.Option) wrapper.Wrapper {
-	w := &WrapImpl{
-		recoveryFunc: grecover.DefaultRecovery,
-	}
-
-	var option = &wrapper.Options{}
+	w := &wrapImpl{}
+	option := &wrapper.Options{}
 	for _, o := range opts {
 		o(option)
 	}
@@ -33,7 +32,7 @@ func New(opts ...wrapper.Option) wrapper.Wrapper {
 }
 
 // Wrap fn func in goroutine to run
-func (w *WrapImpl) Wrap(fn func()) {
+func (w *wrapImpl) Wrap(fn func()) {
 	w.Add(1)
 	go func() {
 		defer w.Done()
@@ -42,7 +41,7 @@ func (w *WrapImpl) Wrap(fn func()) {
 }
 
 // WrapWithRecover exec func with recover
-func (w *WrapImpl) WrapWithRecover(fn func()) {
+func (w *wrapImpl) WrapWithRecover(fn func()) {
 	w.Add(1)
 	go func() {
 		defer w.recoveryFunc()