Merge pull request #4 from anantadwi13/dev

implement func wrapper

Author: anantadwi13

.github/workflows/code-coverage.yml

@@ -3,7 +3,7 @@ name: Test and coverage
 on: [push, pull_request]
 
 env:
-  PACKAGES: io
+  PACKAGES: io, wrapper
 
 jobs:
   test:

wrapper/func.go

@@ -0,0 +1,195 @@
+package wrapper
+
+import (
+	"context"
+	"errors"
+	"reflect"
+	"sync"
+	"sync/atomic"
+)
+
+var (
+	ErrAlreadyShutdown = errors.New("already shutdown")
+)
+
+type HandleFunc func(ctx context.Context, wrapperData *Data)
+
+type Option func(wrapperData *Data)
+
+type Middleware func(next HandleFunc) HandleFunc
+
+type FuncManager interface {
+	// Run will run the fn synchronously
+	Run(ctx context.Context, fn HandleFunc, opts ...Option)
+	// RunAsync will run the fn inside goroutine. No need to spawn the goroutine
+	RunAsync(ctx context.Context, fn HandleFunc, opts ...Option)
+	// Wait will wait for the func manager is shutdown
+	Wait() <-chan struct{}
+	// Shutdown will force shutdown when the ctx is done
+	Shutdown(ctx context.Context) error
+}
+
+type Data struct {
+	dataLock sync.RWMutex
+	data     map[interface{}]interface{}
+}
+
+func (d *Data) Get(key interface{}) interface{} {
+	d.dataLock.RLock()
+	defer d.dataLock.RUnlock()
+	if d.data == nil {
+		return nil
+	}
+	return d.data[key]
+}
+
+func (d *Data) Set(key interface{}, val interface{}) error {
+	if key == nil {
+		return errors.New("nil key")
+	}
+	if !reflect.TypeOf(key).Comparable() {
+		return errors.New("key is not comparable")
+	}
+	d.dataLock.Lock()
+	defer d.dataLock.Unlock()
+	if d.data == nil {
+		d.data = make(map[interface{}]interface{})
+	}
+	d.data[key] = val
+	return nil
+}
+
+type key string
+
+const (
+	keyIdentifier = key("identifier")
+)
+
+func WithOptionIdentifier(funcName string) Option {
+	return func(data *Data) {
+		_ = data.Set(keyIdentifier, funcName)
+	}
+}
+
+func GetIdentifier(wrapperData *Data) string {
+	val, ok := wrapperData.Get(keyIdentifier).(string)
+	if !ok {
+		return ""
+	}
+	return val
+}
+
+func WithMiddlewareRecoverPanic(onPanic func(recoverVal interface{}, wrapperData *Data)) Middleware {
+	return func(next HandleFunc) HandleFunc {
+		return func(ctx context.Context, wrapperData *Data) {
+			defer func() {
+				val := recover()
+				if val != nil {
+					if onPanic != nil {
+						onPanic(val, wrapperData)
+					}
+				}
+			}()
+			next(ctx, wrapperData)
+		}
+	}
+}
+
+type funcManager struct {
+	wg            sync.WaitGroup
+	isShutdown    int32
+	shutdown      chan struct{}
+	mainCtx       context.Context
+	mainCtxCancel context.CancelFunc
+	middlewares   []Middleware
+}
+
+func NewFuncManager(middlewares ...Middleware) FuncManager {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	m := &funcManager{
+		shutdown:      make(chan struct{}),
+		mainCtx:       ctx,
+		mainCtxCancel: cancel,
+		middlewares:   middlewares,
+	}
+
+	return m
+}
+
+func (m *funcManager) Run(ctx context.Context, fn HandleFunc, opts ...Option) {
+	if atomic.LoadInt32(&m.isShutdown) == 1 {
+		return
+	}
+
+	m.wg.Add(1)
+	defer m.wg.Done()
+	m.run(ctx, fn, opts...)
+}
+
+func (m *funcManager) RunAsync(ctx context.Context, fn HandleFunc, opts ...Option) {
+	if atomic.LoadInt32(&m.isShutdown) == 1 {
+		return
+	}
+
+	m.wg.Add(1)
+	go func() {
+		defer m.wg.Done()
+		m.run(ctx, fn, opts...)
+	}()
+}
+
+func (m *funcManager) Wait() <-chan struct{} {
+	return m.shutdown
+}
+
+func (m *funcManager) Shutdown(ctx context.Context) error {
+	if !atomic.CompareAndSwapInt32(&m.isShutdown, 0, 1) {
+		return ErrAlreadyShutdown
+	}
+
+	defer func() {
+		close(m.shutdown)
+	}()
+
+	m.mainCtxCancel()
+
+	done := make(chan struct{})
+	go func() {
+		m.wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+	}
+
+	return nil
+}
+
+func (m *funcManager) run(ctx context.Context, fn HandleFunc, opts ...Option) {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	wrapperData := &Data{}
+
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-m.mainCtx.Done():
+			cancel()
+		}
+	}()
+
+	for _, opt := range opts {
+		opt(wrapperData)
+	}
+
+	for i := len(m.middlewares) - 1; i >= 0; i-- {
+		fn = m.middlewares[i](fn)
+	}
+
+	fn(ctx, wrapperData)
+}

wrapper/func_test.go

@@ -0,0 +1,177 @@
+package wrapper
+
+import (
+	"context"
+	"errors"
+	"log"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestFlow1(t *testing.T) {
+	checker := int32(14)
+	wg := sync.WaitGroup{}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	wg.Add(3) // we will run 3 functions
+	m := NewFuncManager(
+		func(next HandleFunc) HandleFunc {
+			return func(ctx context.Context, wrapperData *Data) {
+				defer wg.Done()
+
+				switch GetIdentifier(wrapperData) {
+				case "check-shutdown", "check-panic", "":
+					atomic.AddInt32(&checker, -1)
+				}
+				atomic.AddInt32(&checker, -1)
+				log.Println("previous", GetIdentifier(wrapperData))
+				next(ctx, wrapperData)
+				atomic.AddInt32(&checker, -1)
+				log.Println("after", GetIdentifier(wrapperData))
+			}
+		},
+		WithMiddlewareRecoverPanic(func(recoverVal interface{}, wrapperData *Data) {
+			atomic.AddInt32(&checker, -1)
+			log.Println("panic", GetIdentifier(wrapperData), recoverVal)
+		}),
+	)
+
+	runnerCtx, runnerCtxCancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer runnerCtxCancel()
+
+	m.RunAsync(runnerCtx, func(ctx context.Context, wrapperData *Data) {
+		<-ctx.Done()
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			atomic.AddInt32(&checker, -1)
+		}
+		panic("trigger panic" + ctx.Err().Error())
+	}, WithOptionIdentifier("check-panic"))
+
+	m.RunAsync(
+		context.Background(),
+		func(ctx context.Context, wrapperData *Data) {
+			select {
+			case <-time.After(20 * time.Second):
+				panic("will not be triggered")
+			case <-ctx.Done():
+				atomic.AddInt32(&checker, -1)
+				log.Println("will be here")
+			}
+		},
+		WithOptionIdentifier("check-shutdown"),
+	)
+
+	m.Run(
+		context.Background(),
+		func(ctx context.Context, wrapperData *Data) {
+			atomic.AddInt32(&checker, -1)
+			log.Println("executed")
+		},
+	)
+
+	go func() {
+		wg.Wait()
+		cancel()
+	}()
+
+	go func() {
+		<-time.After(2 * time.Second)
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		err := m.Shutdown(ctx)
+		if err != nil {
+			log.Println("shutdown error", err)
+			return
+		}
+	}()
+
+	<-m.Wait()
+	<-ctx.Done()
+
+	// running function after shutting down
+	m.Run(context.Background(), func(ctx context.Context, wrapperData *Data) {
+		atomic.AddInt32(&checker, -1) // will not be triggered
+	})
+	m.RunAsync(context.Background(), func(ctx context.Context, wrapperData *Data) {
+		// todo
+		atomic.AddInt32(&checker, -1) // will not be triggered
+	})
+	err := m.Shutdown(context.Background())
+	if errors.Is(err, ErrAlreadyShutdown) {
+		atomic.AddInt32(&checker, -1)
+	}
+
+	if checker != 0 {
+		t.Errorf("invalid checker, checker is not 0. checker: %d", checker)
+	}
+}
+
+func TestFlow2(t *testing.T) {
+	checker := int32(2)
+	wg := sync.WaitGroup{}
+	m := NewFuncManager()
+
+	wg.Add(1)
+	m.RunAsync(context.Background(), func(ctx context.Context, wrapperData *Data) {
+		defer wg.Done()
+		<-time.After(2 * time.Second)
+		atomic.AddInt32(&checker, -1)
+	})
+
+	ctxShutdown, cancelCtxShutdown := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancelCtxShutdown()
+	err := m.Shutdown(ctxShutdown)
+	if errors.Is(err, context.DeadlineExceeded) {
+		atomic.AddInt32(&checker, -1)
+	}
+
+	wg.Wait()
+
+	if checker != 0 {
+		t.Errorf("invalid checker, checker is not 0. checker: %d", checker)
+	}
+}
+
+func TestData(t *testing.T) {
+	checker := int32(6)
+	data := &Data{}
+
+	// key is nil
+	err := data.Set(nil, 123)
+	if err != nil {
+		checker--
+	}
+
+	// non comparable
+	err = data.Set([]string{"a"}, 123)
+	if err != nil {
+		checker--
+	}
+
+	// key as string
+	err = data.Set("abc", 123)
+	if err == nil {
+		checker--
+	}
+
+	// key as another string type
+	type key string
+	err = data.Set(key("abc"), 456)
+	if err == nil {
+		checker--
+	}
+
+	if data.Get("abc") == 123 {
+		checker--
+	}
+	if data.Get(key("abc")) == 456 {
+		checker--
+	}
+
+	if checker != 0 {
+		t.Errorf("invalid checker, checker is not 0. checker: %d", checker)
+	}
+}

wrapper/go.mod

@@ -0,0 +1,3 @@
+module github.com/anantadwi13/go-sdk/wrapper
+
+go 1.14