kpt-grpc-demo/util/cleanup/entry.go

// Package cleanup provides a single point for registering clean up functions.
// This is similar to Finalizer, except that the clean up functions are
// guaranteed to be called if the process terminates normally.
//
// Usage:
//
// In my_package.go
//
//   cleanup.Register(func(){
//     // Arbitrary clean up function, most likely close goroutine, etc.
//   })
//
// In main.go
//
//   func main() {
//     flag.Parse()
//     defer cleanup.Run()
//   }
package cleanup

import (
	"context"
	"reflect"
	"sync"

	log "kpt-grpc-demo/util/logger"
	"kpt-grpc-demo/util/xerr"
)

// entry global cleanup entry
var entry Entry

// Register adds a function to the global cleanup queue.
func Register(container interface{}) {
	entry.Register(container)
}

// Run runs all the global cleanup functions registered.
func Run() {
	entry.Run()
}

type Entry struct {
	mu   sync.Mutex
	fns  []func()
	once sync.Once
}

// Run runs all the cleanup functions registered.
func (entry *Entry) Run() {
	log.Infof("cleanup: performing %d cleanups", len(entry.fns))

	entry.once.Do(func() {
		for _, f := range entry.fns {
			f()
		}
	})

	log.Infof("cleanup: all done")
}

// Register adds a function to the cleanup queue.
func (entry *Entry) Register(container interface{}) {
	v := reflect.Indirect(reflect.ValueOf(container))
	var err error
	switch v.Kind() {
	case reflect.Func:
		err = entry.RegisterFunc(v.Interface())
	case reflect.Struct:
		err = entry.RegisterStruct(v.Interface())
	default:
		panic("cleanup: unsupported type")
	}
	if err != nil {
		panic(err)
	}
}

func (entry *Entry) RegisterStruct(ctor interface{}) error {
	cValue := reflect.Indirect(reflect.ValueOf(ctor))
	if cValue.Kind() != reflect.Struct {
		return xerr.New("RegisterStruct receive a struct or ptr to struct")
	}
	for i := 0; i < cValue.NumField(); i++ {
		field := cValue.Field(i)
		if field.IsZero() {
			continue
		}
		method := field.MethodByName("Close")
		if !method.IsValid() {
			method = field.MethodByName("Flush")
		}
		if method.IsValid() {
			if err := entry.RegisterFunc(method.Interface()); err != nil {
				log.WithError(err).WithField("fieldName", field.Type().Name()).Error("register func failed")
			}
		}
	}
	return nil
}

// RegisterFunc receive func() or func() error
func (entry *Entry) RegisterFunc(fn interface{}) error {
	fType := reflect.TypeOf(fn)
	if fType.Kind() != reflect.Func {
		return xerr.New("cleanup: unsupported type")
	}
	if fType.NumIn() > 0 {
		return xerr.New("RegisterFunc receive func() or func() error")
	}

	if f, ok := fn.(func()); ok {
		entry.register(f)
		return nil
	}
	if f, ok := fn.(func() error); ok {
		entry.register(func() {
			if err := f(); err != nil {
				_ = xerr.ReportError(context.Background(), err)
			}
		})
		return nil
	}
	return xerr.New("RegisterFunc receive func() or func() error")
}

func (entry *Entry) register(fns ...func()) {
	entry.mu.Lock()
	defer entry.mu.Unlock()
	for _, fn := range fns {
		entry.fns = append(entry.fns, fn)
	}
}