buildr/modules/plugin/module.go

package plugin

import (
	"context"
	"errors"
	"fmt"
	"time"

	"connectrpc.com/connect"

	commonv1 "code.icb4dc0.de/buildr/api/generated/common/v1"
	remotev1 "code.icb4dc0.de/buildr/api/generated/remote/v1"
	"code.icb4dc0.de/buildr/api/generated/wasi/v1/rpcv1connect"
	"code.icb4dc0.de/buildr/common/wasirpc"

	"code.icb4dc0.de/buildr/buildr/internal/hcl"
	"code.icb4dc0.de/buildr/buildr/modules"

	"github.com/google/uuid"
	hcl2 "github.com/hashicorp/hcl/v2"
	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
)

var (
	_ modules.Module      = (*Module)(nil)
	_ hcl.Marshaler       = (*Module)(nil)
	_ modules.Helper      = (*Module)(nil)
	_ modules.Initializer = (*Module)(nil)
	_ modules.BinaryNamer = (*Module)(nil)
)

const runtimeClosingTimeout = 1 * time.Second

type Module struct {
	PluginPayload  PayloadReader
	ModuleSpec     map[string]any `hcl:",remain"`
	modSpec        *commonv1.ModuleSpec
	PluginType     string
	PluginCategory modules.Category
}

func (m Module) Execute(ctx modules.ExecutionContext) (err error) {
	runtime := m.prepareWASIRuntime(ctx)

	defer func() {
		ctx, cancel := context.WithTimeout(context.Background(), runtimeClosingTimeout)
		defer cancel()
		if closeErr := runtime.Close(ctx); closeErr != nil {
			err = errors.Join(err, fmt.Errorf("failed to close runtime: %w", closeErr))
		}
	}()

	wasiCtx := &wasiPluginExecContext{
		ctx: ctx,
	}

	_, err = runtime.NewHostModuleBuilder("buildr").
		NewFunctionBuilder().WithFunc(NewHostFuncExport(wasiCtx.log).Call).Export("log_msg").
		NewFunctionBuilder().WithFunc(NewHostFuncExport(wasiCtx.getState).Call).Export("get_state").
		NewFunctionBuilder().WithFunc(NewHostFuncExport(wasiCtx.setState).Call).Export("set_state").
		NewFunctionBuilder().WithFunc(NewHostFuncExport(wasiCtx.exec).Call).Export("exec").
		NewFunctionBuilder().WithFunc(NewHostFuncExport(wasiCtx.lookPath).Call).Export("lookPath").
		Instantiate(ctx)

	if err != nil {
		return fmt.Errorf("failed to instantiate host module: %w", err)
	}

	closer, err := wasi_snapshot_preview1.Instantiate(ctx, runtime)
	if err != nil {
		return fmt.Errorf("failed to instantiate WASI infrastructure: %w", err)
	}

	defer func() {
		if closeErr := closer.Close(context.Background()); closeErr != nil {
			err = errors.Join(err, fmt.Errorf("failed to close WASI runtime: %w", closeErr))
		}
	}()

	moduleConfig := wazero.NewModuleConfig().
		WithStdout(ctx.StdOut()).
		WithStderr(ctx.StdErr()).
		WithFS(newMultiDirFS(ctx.WorkingDir(), ctx.OutDir()))

	pluginPayload, err := m.PluginPayload.Bytes()
	if err != nil {
		return err
	}

	mod, err := runtime.InstantiateWithConfig(ctx, pluginPayload, moduleConfig)
	if err != nil {
		return fmt.Errorf("failed to instantiate WASI module: %w", err)
	}

	moduleClient, err := wasirpc.NewModuleClient(mod)
	if err != nil {
		return fmt.Errorf("failed to create module client: %w", err)
	}

	wasiClient := rpcv1connect.NewWasiExecutorServiceClient(moduleClient, fmt.Sprintf("wasi://%s", m.PluginType))

	startTask := &remotev1.StartTaskRequest{
		Reference: &remotev1.TaskReference{
			Id: uuid.NewString(),
			Module: &commonv1.ModuleReference{
				ModuleCategory: m.PluginCategory,
				ModuleType:     m.PluginType,
			},
			Name: ctx.Name(),
		},
		Buildr: &remotev1.Buildr{
			Repo: &remotev1.Buildr_Repo{
				Root: ctx.WorkingDir(),
			},
			OutDir: ctx.OutDir(),
			BinDir: ctx.BinariesDir(),
		},
		Spec: m.modSpec,
	}

	_, err = wasiClient.StartTask(ctx, connect.NewRequest(startTask))

	return err
}

func (m Module) Category() modules.Category {
	return m.PluginCategory
}

func (m Module) Type() string {
	return m.PluginType
}

func (m *Module) SetModuleSpec(spec *commonv1.ModuleSpec) {
	m.modSpec = spec
}

func (m Module) Init(hclCtx *hcl2.EvalContext) (modules.Module, error) {
	m.modSpec = &commonv1.ModuleSpec{
		Category: m.PluginCategory,
		Type:     m.PluginType,
		Values:   make(map[string]*commonv1.ModuleSpec_Value),
	}

	for k, v := range m.ModuleSpec {
		if a, ok := v.(*hcl2.Attribute); ok {
			val, err := a.Expr.Value(hclCtx)
			if err != nil {
				return nil, err
			}
			m.modSpec.Values[k] = unwrap(val)
		}
	}

	return m, nil
}

func (Module) prepareWASIRuntime(ctx context.Context) wazero.Runtime {
	runtimeConfig := wazero.
		NewRuntimeConfig().
		WithCloseOnContextDone(true)

	r := wazero.NewRuntimeWithConfig(ctx, runtimeConfig)

	return r
}