internal/pluginhost/client_guard.go

Rama main · réplica de lectura del repositorio original.

package pluginhost

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
)

type pluginCallFinishKey struct{}

type pluginCallFinishHook struct {
	once    sync.Once
	claimed atomic.Bool
	fn      func()
}

func withPluginCallFinish(ctx context.Context, fn func()) (context.Context, func()) {
	if ctx == nil {
		ctx = context.Background()
	}
	hook := &pluginCallFinishHook{fn: fn}
	return context.WithValue(ctx, pluginCallFinishKey{}, hook), func() {
		if !hook.claimed.Load() {
			hook.once.Do(hook.fn)
		}
	}
}

func claimPluginCallFinish(ctx context.Context) func() {
	if ctx == nil {
		return nil
	}
	hook, ok := ctx.Value(pluginCallFinishKey{}).(*pluginCallFinishHook)
	if !ok || hook == nil || !hook.claimed.CompareAndSwap(false, true) {
		return nil
	}
	return func() {
		hook.once.Do(hook.fn)
	}
}

type guardedPluginClient struct {
	mu           sync.Mutex
	cond         *sync.Cond
	inner        pluginClient
	instance     *hostCallbackInstance
	calls        int
	closed       bool
	shutdownDone chan struct{}
}

func newGuardedPluginClient(inner pluginClient) *guardedPluginClient {
	instance := pluginCallbackInstance(inner)
	if instance == nil {
		instance = &hostCallbackInstance{}
	}
	client := &guardedPluginClient{inner: inner, instance: instance, shutdownDone: make(chan struct{})}
	client.cond = sync.NewCond(&client.mu)
	return client
}

func (c *guardedPluginClient) callbackInstance() *hostCallbackInstance {
	if c == nil {
		return nil
	}
	return c.instance
}

func pluginCallbackInstance(client pluginClient) *hostCallbackInstance {
	if client == nil {
		return nil
	}
	if provider, ok := client.(interface{ callbackInstance() *hostCallbackInstance }); ok {
		return provider.callbackInstance()
	}
	return nil
}

func (c *guardedPluginClient) Call(ctx context.Context, method string, request []byte) ([]byte, error) {
	inner, errAcquire := c.acquire()
	if errAcquire != nil {
		return nil, errAcquire
	}
	if ctx == nil {
		ctx = context.Background()
	}
	onFinish := claimPluginCallFinish(ctx)
	result := make(chan guardedPluginCallResult, 1)
	go func() {
		defer c.release()
		defer func() {
			if recovered := recover(); recovered != nil {
				if onFinish != nil {
					onFinish()
				}
				result <- guardedPluginCallResult{recovered: recovered}
			}
		}()
		response, errCall := inner.Call(ctx, method, request)
		if onFinish != nil {
			onFinish()
		}
		result <- guardedPluginCallResult{response: response, err: errCall}
	}()
	select {
	case callResult := <-result:
		if callResult.recovered != nil {
			panic(callResult.recovered)
		}
		return callResult.response, callResult.err
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

type guardedPluginCallResult struct {
	response  []byte
	err       error
	recovered any
}

func (c *guardedPluginClient) acquire() (pluginClient, error) {
	if c == nil {
		return nil, fmt.Errorf("plugin client is closed")
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed || c.inner == nil {
		return nil, fmt.Errorf("plugin client is closed")
	}
	c.calls++
	return c.inner, nil
}

func (c *guardedPluginClient) release() {
	c.mu.Lock()
	c.calls--
	if c.calls == 0 {
		c.cond.Broadcast()
	}
	c.mu.Unlock()
}

func (c *guardedPluginClient) Shutdown() {
	c.ShutdownContext(context.Background())
}

// ShutdownContext detaches the client immediately and waits for active calls only
// until ctx is canceled. Detached cleanup continues asynchronously when needed.
func (c *guardedPluginClient) ShutdownContext(ctx context.Context) {
	if c == nil {
		return
	}
	if ctx == nil {
		ctx = context.Background()
	}

	c.mu.Lock()
	if c.closed {
		done := c.shutdownDone
		c.mu.Unlock()
		select {
		case <-done:
		case <-ctx.Done():
		}
		return
	}
	c.closed = true
	inner := c.inner
	c.inner = nil
	done := c.shutdownDone
	c.mu.Unlock()

	go func() {
		c.mu.Lock()
		for c.calls > 0 {
			c.cond.Wait()
		}
		c.mu.Unlock()
		if inner != nil {
			inner.Shutdown()
		}
		close(done)
	}()

	select {
	case <-done:
	case <-ctx.Done():
	}
}