registry.go raw
1 package organ
2
3 import (
4 "context"
5 "fmt"
6 "sort"
7 "sync"
8 "time"
9
10 "github.com/tetratelabs/wazero"
11 "github.com/tetratelabs/wazero/api"
12 "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
13 )
14
15 // Registry manages loaded organs for one lattice instance.
16 // It owns the wazero runtime and handles module lifecycle.
17 type Registry struct {
18 mu sync.RWMutex
19 organs map[OrganID]*loadedOrgan
20 byType map[OrganType][]*loadedOrgan
21 runtime wazero.Runtime
22 ctx context.Context
23 }
24
25 // loadedOrgan extends Organ with the wazero module instance.
26 type loadedOrgan struct {
27 Organ
28 compiled wazero.CompiledModule
29 instance api.Module
30 }
31
32 // NewRegistry creates a registry with a fresh wazero runtime.
33 func NewRegistry(ctx context.Context) (*Registry, error) {
34 r := wazero.NewRuntime(ctx)
35
36 // Instantiate WASI for organs that use it.
37 if _, err := wasi_snapshot_preview1.Instantiate(ctx, r); err != nil {
38 r.Close(ctx)
39 return nil, fmt.Errorf("wasi init: %w", err)
40 }
41
42 return &Registry{
43 organs: make(map[OrganID]*loadedOrgan),
44 byType: make(map[OrganType][]*loadedOrgan),
45 runtime: r,
46 ctx: ctx,
47 }, nil
48 }
49
50 // Load compiles and instantiates a WASM module.
51 // The organ must export: alloc, dealloc, organ_type, organ_version, process.
52 func (reg *Registry) Load(wasmBytes []byte, manifest Manifest) (*Organ, error) {
53 reg.mu.Lock()
54 defer reg.mu.Unlock()
55
56 id := ComputeID(wasmBytes)
57 manifest.ID = id
58 manifest.Size = len(wasmBytes)
59
60 // Already loaded?
61 if existing, ok := reg.organs[id]; ok {
62 return &existing.Organ, nil
63 }
64
65 // Compile.
66 compiled, err := reg.runtime.CompileModule(reg.ctx, wasmBytes)
67 if err != nil {
68 return nil, fmt.Errorf("compile: %w", err)
69 }
70
71 // Instantiate as a WASI reactor (calls _initialize, not _start).
72 config := wazero.NewModuleConfig().
73 WithName(fmt.Sprintf("organ-%s", id)).
74 WithStartFunctions("_initialize")
75
76 instance, err := reg.runtime.InstantiateModule(reg.ctx, compiled, config)
77 if err != nil {
78 return nil, fmt.Errorf("instantiate: %w", err)
79 }
80
81 // Verify required exports exist.
82 for _, name := range []string{"alloc", "dealloc", "process"} {
83 if instance.ExportedFunction(name) == nil {
84 instance.Close(reg.ctx)
85 return nil, fmt.Errorf("organ missing required export: %s", name)
86 }
87 }
88
89 lo := &loadedOrgan{
90 Organ: Organ{
91 Manifest: manifest,
92 WasmBytes: wasmBytes,
93 Status: StatusLoaded,
94 Loaded: time.Now(),
95 },
96 compiled: compiled,
97 instance: instance,
98 }
99
100 reg.organs[id] = lo
101 reg.byType[manifest.Type] = append(reg.byType[manifest.Type], lo)
102
103 return &lo.Organ, nil
104 }
105
106 // Call invokes an organ's process function with the given opcode and input.
107 // The opcode byte is prepended to the input before passing to the WASM module.
108 func (reg *Registry) Call(id OrganID, opcode byte, input []byte) ([]byte, error) {
109 reg.mu.RLock()
110 lo, ok := reg.organs[id]
111 reg.mu.RUnlock()
112
113 if !ok {
114 return nil, fmt.Errorf("organ not found: %s", id)
115 }
116 if lo.Status != StatusLoaded {
117 return nil, fmt.Errorf("organ not loaded: %s (status %d)", id, lo.Status)
118 }
119
120 mod := lo.instance
121
122 // Prepend opcode.
123 payload := make([]byte, 1+len(input))
124 payload[0] = opcode
125 copy(payload[1:], input)
126
127 // Allocate in guest memory.
128 allocFn := mod.ExportedFunction("alloc")
129 results, err := allocFn.Call(reg.ctx, uint64(len(payload)))
130 if err != nil {
131 return nil, fmt.Errorf("alloc: %w", err)
132 }
133 ptr := uint32(results[0])
134
135 // Write input to guest memory.
136 if !mod.Memory().Write(ptr, payload) {
137 return nil, fmt.Errorf("memory write failed at ptr=%d len=%d", ptr, len(payload))
138 }
139
140 // Call process.
141 processFn := mod.ExportedFunction("process")
142 results, err = processFn.Call(reg.ctx, uint64(ptr), uint64(len(payload)))
143 if err != nil {
144 return nil, fmt.Errorf("process: %w", err)
145 }
146
147 // Unpack result: high 32 bits = pointer, low 32 bits = length.
148 packed := results[0]
149 outPtr := uint32(packed >> 32)
150 outLen := uint32(packed & 0xFFFFFFFF)
151
152 // Read output from guest memory.
153 output, ok := mod.Memory().Read(outPtr, outLen)
154 if !ok {
155 return nil, fmt.Errorf("memory read failed at ptr=%d len=%d", outPtr, outLen)
156 }
157
158 // Make a copy — wazero memory may be invalidated.
159 result := make([]byte, len(output))
160 copy(result, output)
161
162 // Free both buffers in guest.
163 deallocFn := mod.ExportedFunction("dealloc")
164 deallocFn.Call(reg.ctx, uint64(ptr), uint64(len(payload)))
165 deallocFn.Call(reg.ctx, uint64(outPtr), uint64(outLen))
166
167 lo.CallCount.Add(1)
168 lo.LastCall = time.Now()
169
170 return result, nil
171 }
172
173 // Evict destroys an organ: closes the wazero module, releases memory,
174 // removes from registry.
175 func (reg *Registry) Evict(id OrganID) error {
176 reg.mu.Lock()
177 defer reg.mu.Unlock()
178
179 lo, ok := reg.organs[id]
180 if !ok {
181 return fmt.Errorf("organ not found: %s", id)
182 }
183
184 // Close the WASM instance.
185 if lo.instance != nil {
186 lo.instance.Close(reg.ctx)
187 }
188
189 lo.Status = StatusEvicted
190 lo.WasmBytes = nil // release for GC
191
192 // Remove from type index.
193 organs := reg.byType[lo.Manifest.Type]
194 for i, o := range organs {
195 if o.Manifest.ID == id {
196 reg.byType[lo.Manifest.Type] = append(organs[:i], organs[i+1:]...)
197 break
198 }
199 }
200
201 delete(reg.organs, id)
202 return nil
203 }
204
205 // Best returns the highest-fitness organ of the given type, or nil.
206 func (reg *Registry) Best(t OrganType) *Organ {
207 reg.mu.RLock()
208 defer reg.mu.RUnlock()
209
210 organs := reg.byType[t]
211 if len(organs) == 0 {
212 return nil
213 }
214
215 var best *loadedOrgan
216 for _, o := range organs {
217 if o.Status != StatusLoaded {
218 continue
219 }
220 if best == nil || best.Manifest.Fitness.Less(o.Manifest.Fitness) {
221 best = o
222 }
223 }
224
225 if best == nil {
226 return nil
227 }
228 return &best.Organ
229 }
230
231 // All returns all loaded organs of the given type, sorted by fitness descending.
232 func (reg *Registry) All(t OrganType) []*Organ {
233 reg.mu.RLock()
234 defer reg.mu.RUnlock()
235
236 organs := reg.byType[t]
237 result := make([]*Organ, 0, len(organs))
238 for _, o := range organs {
239 if o.Status == StatusLoaded {
240 result = append(result, &o.Organ)
241 }
242 }
243
244 sort.Slice(result, func(i, j int) bool {
245 return result[j].Manifest.Fitness.Less(result[i].Manifest.Fitness)
246 })
247
248 return result
249 }
250
251 // Get returns an organ by ID, or nil if not found.
252 func (reg *Registry) Get(id OrganID) *Organ {
253 reg.mu.RLock()
254 defer reg.mu.RUnlock()
255
256 if lo, ok := reg.organs[id]; ok {
257 return &lo.Organ
258 }
259 return nil
260 }
261
262 // Close shuts down the registry and all loaded organs.
263 func (reg *Registry) Close() error {
264 reg.mu.Lock()
265 defer reg.mu.Unlock()
266
267 // wazero runtime close destroys all modules.
268 return reg.runtime.Close(reg.ctx)
269 }
270