ollama source for Momentry Core verification

x/internal/mlxthread/thread.go

@@ -0,0 +1,183 @@
+package mlxthread
+
+import (
+	"context"
+	"errors"
+	"runtime"
+	"sync/atomic"
+)
+
+var ErrStopped = errors.New("mlx thread stopped")
+
+type Thread struct {
+	name string
+
+	jobs     chan job
+	done     chan struct{}
+	stopping atomic.Bool
+}
+
+type job struct {
+	fn     func() error
+	result chan result
+	stop   bool
+}
+
+type result struct {
+	err        error
+	panicValue any
+}
+
+// Start creates a long-lived worker goroutine locked to one OS thread.
+func Start(name string, init func() error) (*Thread, error) {
+	t := &Thread{
+		name: name,
+		jobs: make(chan job),
+		done: make(chan struct{}),
+	}
+
+	initResult := make(chan result, 1)
+	go t.loop(init, initResult)
+
+	res := <-initResult
+	if res.panicValue != nil {
+		panic(res.panicValue)
+	}
+	if res.err != nil {
+		return nil, res.err
+	}
+
+	return t, nil
+}
+
+// Do runs fn on the locked OS thread.
+//
+// Context cancellation only applies while the work is queued. Once the worker
+// accepts a job, the job runs until fn returns or reaches its own cancellation
+// checks.
+func (t *Thread) Do(ctx context.Context, fn func() error) error {
+	res, err := t.enqueue(ctx, fn, false, false)
+	if err != nil {
+		return err
+	}
+	if res.panicValue != nil {
+		panic(res.panicValue)
+	}
+	return res.err
+}
+
+func Call[T any](ctx context.Context, t *Thread, fn func() (T, error)) (T, error) {
+	var value T
+	err := t.Do(ctx, func() error {
+		var err error
+		value, err = fn()
+		return err
+	})
+	return value, err
+}
+
+// Stop runs cleanup on the locked OS thread and then shuts the worker down.
+func (t *Thread) Stop(ctx context.Context, cleanup func()) error {
+	ctx = contextOrBackground(ctx)
+
+	if !t.stopping.CompareAndSwap(false, true) {
+		select {
+		case <-t.done:
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+
+	res, err := t.enqueue(ctx, func() error {
+		if cleanup != nil {
+			cleanup()
+		}
+		return nil
+	}, true, true)
+	if err != nil {
+		if !errors.Is(err, ErrStopped) {
+			t.stopping.Store(false)
+		}
+		return err
+	}
+	if res.panicValue != nil {
+		panic(res.panicValue)
+	}
+	if res.err != nil {
+		return res.err
+	}
+
+	select {
+	case <-t.done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func (t *Thread) loop(init func() error, initResult chan<- result) {
+	runtime.LockOSThread()
+	// Deliberately do not unlock. MLX thread-local state belongs to this worker
+	// until shutdown so it cannot leak back to arbitrary Go goroutines.
+
+	res := run(init)
+	initResult <- res
+	if res.err != nil || res.panicValue != nil {
+		close(t.done)
+		return
+	}
+
+	for {
+		j := <-t.jobs
+		res := run(j.fn)
+		j.result <- res
+		if j.stop {
+			close(t.done)
+			return
+		}
+	}
+}
+
+func (t *Thread) enqueue(ctx context.Context, fn func() error, stop, allowStopping bool) (result, error) {
+	ctx = contextOrBackground(ctx)
+	if err := ctx.Err(); err != nil {
+		return result{}, err
+	}
+
+	if !allowStopping && t.stopping.Load() {
+		return result{}, ErrStopped
+	}
+
+	resultCh := make(chan result, 1)
+	j := job{fn: fn, result: resultCh, stop: stop}
+
+	select {
+	case <-ctx.Done():
+		return result{}, ctx.Err()
+	case <-t.done:
+		return result{}, ErrStopped
+	case t.jobs <- j:
+	}
+
+	return <-resultCh, nil
+}
+
+func run(fn func() error) (res result) {
+	defer func() {
+		if v := recover(); v != nil {
+			res.panicValue = v
+		}
+	}()
+	if fn != nil {
+		res.err = fn()
+	}
+	return res
+}
+
+func contextOrBackground(ctx context.Context) context.Context {
+	if ctx != nil {
+		return ctx
+	}
+	return context.Background()
+}

x/internal/mlxthread/thread_affinity_test.go

@@ -0,0 +1,32 @@
+//go:build darwin || linux
+
+package mlxthread
+
+import (
+	"context"
+	"fmt"
+	"testing"
+)
+
+func TestDoUsesSameOSThread(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer thread.Stop(context.Background(), nil)
+
+	var first uint64
+	for range 32 {
+		if err := thread.Do(context.Background(), func() error {
+			id := currentThreadID()
+			if first == 0 {
+				first = id
+			} else if id != first {
+				return fmt.Errorf("job ran on OS thread %d, want %d", id, first)
+			}
+			return nil
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}
+}

x/internal/mlxthread/thread_test.go

@@ -0,0 +1,351 @@
+package mlxthread
+
+import (
+	"context"
+	"errors"
+	"reflect"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestDoRunsInOrder(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer thread.Stop(context.Background(), nil)
+
+	var got []int
+	for i := 0; i < 5; i++ {
+		i := i
+		if err := thread.Do(context.Background(), func() error {
+			got = append(got, i)
+			return nil
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	if want := []int{0, 1, 2, 3, 4}; !reflect.DeepEqual(got, want) {
+		t.Fatalf("got %v, want %v", got, want)
+	}
+}
+
+func TestDoPropagatesPanicToCaller(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer thread.Stop(context.Background(), nil)
+
+	defer func() {
+		if got := recover(); got != "boom" {
+			t.Fatalf("got panic %v, want boom", got)
+		}
+	}()
+
+	_ = thread.Do(context.Background(), func() error {
+		panic("boom")
+	})
+}
+
+func TestDoCancelsBeforeJobStarts(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer thread.Stop(context.Background(), nil)
+
+	running := make(chan struct{})
+	release := make(chan struct{})
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- thread.Do(context.Background(), func() error {
+			close(running)
+			<-release
+			return nil
+		})
+	}()
+	<-running
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err = thread.Do(ctx, func() error {
+		t.Fatal("canceled job should not run")
+		return nil
+	})
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("got %v, want %v", err, context.Canceled)
+	}
+
+	close(release)
+	if err := <-errCh; err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestAlreadyCanceledContextDoesNotEnqueue(t *testing.T) {
+	t.Run("Do", func(t *testing.T) {
+		thread, err := Start("test", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer thread.Stop(context.Background(), nil)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		ran := false
+		err = thread.Do(ctx, func() error {
+			ran = true
+			return nil
+		})
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("got %v, want %v", err, context.Canceled)
+		}
+		if ran {
+			t.Fatal("canceled job ran")
+		}
+	})
+
+	t.Run("Stop", func(t *testing.T) {
+		thread, err := Start("test", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer thread.Stop(context.Background(), nil)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		cleaned := false
+		err = thread.Stop(ctx, func() {
+			cleaned = true
+		})
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("got %v, want %v", err, context.Canceled)
+		}
+		if cleaned {
+			t.Fatal("cleanup ran for canceled stop")
+		}
+		if err := thread.Do(context.Background(), func() error { return nil }); err != nil {
+			t.Fatalf("thread did not accept work after canceled Stop: %v", err)
+		}
+	})
+}
+
+func TestCallReturnsValue(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer thread.Stop(context.Background(), nil)
+
+	got, err := Call(context.Background(), thread, func() (int, error) {
+		return 42, nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got != 42 {
+		t.Fatalf("got %d, want 42", got)
+	}
+}
+
+func TestDoRunsConcurrentlySubmittedWorkSerially(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer thread.Stop(context.Background(), nil)
+
+	oldProcs := runtime.GOMAXPROCS(8)
+	defer runtime.GOMAXPROCS(oldProcs)
+
+	const goroutines = 16
+	const iterations = 64
+
+	var active atomic.Int32
+	var count atomic.Int64
+	var wg sync.WaitGroup
+	errCh := make(chan error, goroutines)
+
+	for range goroutines {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			for range iterations {
+				if err := thread.Do(context.Background(), func() error {
+					if got := active.Add(1); got != 1 {
+						return errors.New("thread executed jobs concurrently")
+					}
+					runtime.Gosched()
+					count.Add(1)
+					if got := active.Add(-1); got != 0 {
+						return errors.New("thread active count did not return to zero")
+					}
+					return nil
+				}); err != nil {
+					errCh <- err
+					return
+				}
+			}
+		}()
+	}
+
+	wg.Wait()
+	close(errCh)
+
+	for err := range errCh {
+		t.Fatal(err)
+	}
+	if got, want := count.Load(), int64(goroutines*iterations); got != want {
+		t.Fatalf("got %d jobs, want %d", got, want)
+	}
+}
+
+func TestStopRunsCleanupAndRejectsWork(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cleaned := 0
+	if err := thread.Stop(context.Background(), func() {
+		cleaned++
+	}); err != nil {
+		t.Fatal(err)
+	}
+	if cleaned != 1 {
+		t.Fatalf("cleanup ran %d times, want 1", cleaned)
+	}
+
+	if err := thread.Stop(context.Background(), func() {
+		cleaned++
+	}); err != nil {
+		t.Fatal(err)
+	}
+	if cleaned != 1 {
+		t.Fatalf("cleanup ran %d times after second Stop, want 1", cleaned)
+	}
+
+	err = thread.Do(context.Background(), func() error {
+		t.Fatal("job should not run after stop")
+		return nil
+	})
+	if !errors.Is(err, ErrStopped) {
+		t.Fatalf("got %v, want %v", err, ErrStopped)
+	}
+}
+
+func TestStopCanceledBeforeEnqueueCanBeRetried(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer thread.Stop(context.Background(), nil)
+
+	running := make(chan struct{})
+	release := make(chan struct{})
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- thread.Do(context.Background(), func() error {
+			close(running)
+			<-release
+			return nil
+		})
+	}()
+	<-running
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+
+	cleanupRan := false
+	err = thread.Stop(ctx, func() {
+		cleanupRan = true
+	})
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("got %v, want %v", err, context.DeadlineExceeded)
+	}
+	if cleanupRan {
+		t.Fatal("cleanup ran even though stop was not enqueued")
+	}
+
+	close(release)
+	if err := <-errCh; err != nil {
+		t.Fatal(err)
+	}
+
+	if err := thread.Do(context.Background(), func() error { return nil }); err != nil {
+		t.Fatalf("thread did not accept work after canceled Stop: %v", err)
+	}
+
+	cleanupRan = false
+	if err := thread.Stop(context.Background(), func() {
+		cleanupRan = true
+	}); err != nil {
+		t.Fatal(err)
+	}
+	if !cleanupRan {
+		t.Fatal("cleanup did not run on retried Stop")
+	}
+}
+
+func TestStopWaitsForActiveWorkBeforeCleanup(t *testing.T) {
+	thread, err := Start("test", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	running := make(chan struct{})
+	release := make(chan struct{})
+	jobErr := make(chan error, 1)
+	go func() {
+		jobErr <- thread.Do(context.Background(), func() error {
+			close(running)
+			<-release
+			return nil
+		})
+	}()
+	<-running
+
+	cleaned := make(chan struct{})
+	stopErr := make(chan error, 1)
+	go func() {
+		stopErr <- thread.Stop(context.Background(), func() {
+			close(cleaned)
+		})
+	}()
+
+	select {
+	case <-cleaned:
+		t.Fatal("cleanup ran before active job completed")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	err = thread.Do(context.Background(), func() error {
+		return errors.New("work should be rejected once Stop starts")
+	})
+	if !errors.Is(err, ErrStopped) {
+		t.Fatalf("got %v, want %v", err, ErrStopped)
+	}
+
+	close(release)
+	if err := <-jobErr; err != nil {
+		t.Fatal(err)
+	}
+	if err := <-stopErr; err != nil {
+		t.Fatal(err)
+	}
+
+	select {
+	case <-cleaned:
+	default:
+		t.Fatal("cleanup did not run")
+	}
+}

x/internal/mlxthread/threadid_darwin_test.go

@@ -0,0 +1,10 @@
+//go:build darwin
+
+package mlxthread
+
+import "syscall"
+
+func currentThreadID() uint64 {
+	id, _, _ := syscall.RawSyscall(syscall.SYS_THREAD_SELFID, 0, 0, 0)
+	return uint64(id)
+}

x/internal/mlxthread/threadid_linux_test.go

@@ -0,0 +1,9 @@
+//go:build linux
+
+package mlxthread
+
+import "syscall"
+
+func currentThreadID() uint64 {
+	return uint64(syscall.Gettid())
+}