components/registry.go

package components

import (
	"context"
	"sort"

	"github.com/urfave/cli/v2"
	"go-micro.dev/v4"
	"go-micro.dev/v4/errors"
	"go-micro.dev/v4/server"
)

var (
	errorRetrievingComponents = errors.InternalServerError("RETRIEVING_COMPONENTS", "error while retrieving components")
	errorComponentsIsNil      = errors.InternalServerError("COMPONENTS_NIL", "components is nil")
)

type HealthInfo struct {
	Message string `json:"message"`
	IsError bool   `json:"is_error"`
}

type HealthInfoMap map[string]HealthInfo

type RegistryKey struct{}

type Registry struct {
	components map[string]Component

	service    micro.Service
	flagPrefix string
}

func Context(ctx context.Context) (*Registry, error) {
	c, ok := ctx.Value(RegistryKey{}).(*Registry)
	if !ok {
		return nil, errorRetrievingComponents
	}

	if c == nil {
		return nil, errorComponentsIsNil
	}

	return c, nil
}

func Must(ctx context.Context) *Registry {
	c, err := Context(ctx)
	if err != nil {
		panic(err)
	}

	return c
}

func New(service micro.Service, flagPrefix string, components ...Component) *Registry {
	com := &Registry{service: service, flagPrefix: flagPrefix, components: make(map[string]Component)}

	com.Add(components...)

	return com
}

func (c *Registry) Service() micro.Service {
	return c.service
}

func (c *Registry) FlagPrefix() string {
	return c.flagPrefix
}

func (c *Registry) Add(components ...Component) {
	for _, component := range components {
		if component == nil {
			continue
		}

		if _, ok := c.components[component.Name()]; ok {
			continue
		}

		c.components[component.Name()] = component
	}
}

func (c *Registry) Get(name string) (Component, error) {
	if component, ok := c.components[name]; ok {
		return component, nil
	}

	return nil, errors.InternalServerError("COMPONENT_NOT_FOUND", "component '%s' has not been found", name)
}

func (c *Registry) Must(name string) Component {
	component, err := c.Get(name)
	if err != nil {
		panic(err)
	}

	return component
}

func (c *Registry) AppendFlags(flags []cli.Flag) []cli.Flag {
	for _, component := range c.components {
		flags = append(flags, component.Flags(c)...)
	}

	return flags
}

func (c *Registry) Initialized() bool {
	for _, component := range c.components {
		if !component.Initialized() {
			return false
		}
	}

	return true
}

func (c *Registry) Init(context *cli.Context) error {
	// Sort Components by Priority ASC
	components := make([]Component, len(c.components))
	for _, com := range c.components {
		components = append(components, com)
	}
	sort.SliceStable(components, func(i, j int) bool {
		if components[i] == nil || components[j] == nil {
			return false
		}

		return components[i].Priority() < components[j].Priority()
	})

	// Init them sorted now
	for _, component := range components {
		if component == nil {
			continue
		}

		if err := component.Init(c, context); err != nil {
			return err
		}
	}

	return nil
}

func (c *Registry) Stop() error {
	// Sort Components by Priority DESC
	components := make([]Component, len(c.components))
	for _, com := range c.components {
		components = append(components, com)
	}
	sort.SliceStable(components, func(i, j int) bool {
		if components[i] == nil || components[j] == nil {
			return false
		}

		return components[i].Priority() > components[j].Priority()
	})

	// Init them sorted now
	for _, component := range components {
		if component == nil {
			continue
		}

		if err := component.Stop(); err != nil {
			return err
		}
	}

	return nil
}

func (c *Registry) Health(context context.Context) HealthInfoMap {
	result := make(HealthInfoMap, len(c.components))

	for _, component := range c.components {
		err := component.Health(context)
		m := "All fine"
		isError := false
		if err != nil {
			m = err.Error()
			isError = true
		}
		result[component.Name()] = HealthInfo{Message: m, IsError: isError}
	}

	return result
}

func (c *Registry) WrapHandler() server.HandlerWrapper {
	return func(h server.HandlerFunc) server.HandlerFunc {
		return func(ctx context.Context, req server.Request, rsp interface{}) error {
			ctx = context.WithValue(ctx, RegistryKey{}, c)

			for _, component := range c.components {
				if err := component.WrapHandlerFunc(ctx, req, rsp); err != nil {
					return err
				}
			}

			return h(ctx, req, rsp)
		}
	}
}