feat(components): define registry interface for components init simplification (#650)

Address #511.

---------

Signed-off-by: Gyuho Lee <[email protected]>

Author: gyuho

components/registry.go

@@ -0,0 +1,80 @@
+package components
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/leptonai/gpud/pkg/eventstore"
+	pkghost "github.com/leptonai/gpud/pkg/host"
+	nvidianvml "github.com/leptonai/gpud/pkg/nvidia-query/nvml"
+)
+
+// GPUdInstance is the instance of the GPUd dependencies.
+type GPUdInstance struct {
+	RootCtx context.Context
+
+	NVMLInstance     nvidianvml.InstanceV2
+	EventStore       eventstore.Store
+	RebootEventStore pkghost.RebootEventStore
+}
+
+// Registry is the interface for the registry of components.
+type Registry interface {
+	// MustRegister registers a component with the given name and initialization function.
+	// It panics if the component is already registered.
+	// It panics if the initialization function returns an error.
+	MustRegister(name string, initFunc func(GPUdInstance) (Component, error))
+}
+
+var _ Registry = &registry{}
+
+type registry struct {
+	mu           sync.RWMutex
+	gpudInstance GPUdInstance
+	components   map[string]Component
+}
+
+// NewRegistry creates a new registry.
+func NewRegistry(gpudInstance GPUdInstance) *registry {
+	return &registry{
+		gpudInstance: gpudInstance,
+		components:   make(map[string]Component),
+	}
+}
+
+// MustRegister registers a component with the given name and initialization function.
+// It panics if the component is already registered.
+// It panics if the initialization function returns an error.
+func (r *registry) MustRegister(name string, initFunc func(GPUdInstance) (Component, error)) {
+	if err := r.registerInit(name, initFunc); err != nil {
+		panic(err)
+	}
+}
+
+// hasRegistered checks if a component with the given name is already registered.
+func (r *registry) hasRegistered(name string) bool {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	_, ok := r.components[name]
+	return ok
+}
+
+// registerInit registers an initialization function for a component with the given name.
+func (r *registry) registerInit(name string, initFunc func(GPUdInstance) (Component, error)) error {
+	if r.hasRegistered(name) {
+		return fmt.Errorf("component %s already registered", name)
+	}
+
+	c, err := initFunc(r.gpudInstance)
+	if err != nil {
+		return err
+	}
+
+	r.mu.Lock()
+	r.components[name] = c
+	r.mu.Unlock()
+
+	return nil
+}

components/registry_test.go

@@ -0,0 +1,66 @@
+package components
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// Mock function that returns a component without error
+func mockInitFuncSuccess(instance GPUdInstance) (Component, error) {
+	return newMockComponent("test-component"), nil
+}
+
+// Mock function that returns an error
+func mockInitFuncError(instance GPUdInstance) (Component, error) {
+	return nil, fmt.Errorf("mock init error")
+}
+
+func TestHasRegistered(t *testing.T) {
+	// Create a new registry
+	reg := NewRegistry(GPUdInstance{
+		RootCtx: context.Background(),
+	})
+
+	// When registry is empty, should return false
+	assert.False(t, reg.hasRegistered("test-component"))
+
+	// Register a component
+	comp := newMockComponent("test-component")
+	reg.mu.Lock()
+	reg.components["test-component"] = comp
+	reg.mu.Unlock()
+
+	// Should now return true for the registered component
+	assert.True(t, reg.hasRegistered("test-component"))
+
+	// Should still return false for unregistered components
+	assert.False(t, reg.hasRegistered("unknown-component"))
+}
+
+func TestRegisterInitFunc(t *testing.T) {
+	// Create a new registry
+	reg := NewRegistry(GPUdInstance{
+		RootCtx: context.Background(),
+	})
+
+	// Test registering a component successfully
+	err := reg.registerInit("test-component", mockInitFuncSuccess)
+	assert.NoError(t, err)
+	assert.True(t, reg.hasRegistered("test-component"))
+
+	// Test registering a component that already exists
+	err = reg.registerInit("test-component", mockInitFuncSuccess)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "already registered")
+
+	// Test registering a component with an initialization function that returns an error
+	err = reg.registerInit("error-component", mockInitFuncError)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "mock init error")
+
+	// The component should not be registered if the init function fails
+	assert.False(t, reg.hasRegistered("error-component"))
+}