Repository: grafana/grafana-plugin-sdk-go (go)

Issue:
Instancemgmt: Instance possibly being disposed while still being in use
Follow up from https://github.com/grafana/grafana/pull/75171 and https://github.com/grafana/grafana/pull/75249 we have concluded that app/datasource instance caching works as expected given a high rate of concurrent requests. However, it seems like it's possible that `Dispose()` of a cached instance being called in a goroutine while there's another goroutine still using it. This could happen in the following way:

1. Goroutine ONE enters this code for cachekey `ABC` and get a hold of a cached instance:
https://github.com/grafana/grafana-plugin-sdk-go/blob/dd41f8d709969a5efd5375833f9754cecf9176bc/backend/instancemgmt/instance_manager.go#L92-L108
2. Goroutine TWO enters this code almost at the same time for the same cachekey `ABC`, but compared to 1) the instance needs an update and goes into this code which will call `Dispose()` of the currently active instance:
https://github.com/grafana/grafana-plugin-sdk-go/blob/dd41f8d709969a5efd5375833f9754cecf9176bc/backend/instancemgmt/instance_manager.go#L110-L137
3. Goroutine ONE continue its execution and tries to use the instance, but fails since it's already disposed.

This should theoretically only happen with a high rate of concurrent requests. 

**Proposed solution:**
We're going to look into disposing instances after X amount of time has passed and/or when the instance no longer is in use.

Pull request:
Instance management: Dispose of instance after 5s
**What this PR does / why we need it**:
Alternative to #762 and #763. This currently doesn't work because haven't put any effort to resolve the tests, but mainly wanted to show the changes in instance_manager.go

**Which issue(s) this PR fixes**:

Fixes #753

**Special notes for your reviewer**:


Patch:
diff --git a/backend/data_adapter.go b/backend/data_adapter.go
index 378bc09..8f54bdf 100644
--- a/backend/data_adapter.go
+++ b/backend/data_adapter.go
@@ -48,6 +48,7 @@ func (a *dataSDKAdapter) QueryData(ctx context.Context, req *pluginv2.QueryDataR
 	ctx = WithGrafanaConfig(ctx, NewGrafanaCfg(req.PluginContext.GrafanaConfig))
 	parsedRequest := FromProto().QueryDataRequest(req)
 	ctx = withHeaderMiddleware(ctx, parsedRequest.GetHTTPHeaders())
+	ctx = withContextualLogAttributes(ctx, parsedRequest.PluginContext, endpointQueryData)
 	resp, err := a.queryDataHandler.QueryData(ctx, parsedRequest)
 	if err != nil {
 		return nil, err
diff --git a/backend/diagnostics_adapter.go b/backend/diagnostics_adapter.go
index 13b3545..8f2ee89 100644
--- a/backend/diagnostics_adapter.go
+++ b/backend/diagnostics_adapter.go
@@ -50,6 +50,7 @@ func (a *diagnosticsSDKAdapter) CheckHealth(ctx context.Context, protoReq *plugi
 		ctx = WithGrafanaConfig(ctx, NewGrafanaCfg(protoReq.PluginContext.GrafanaConfig))
 		parsedReq := FromProto().CheckHealthRequest(protoReq)
 		ctx = withHeaderMiddleware(ctx, parsedReq.GetHTTPHeaders())
+		ctx = withContextualLogAttributes(ctx, parsedReq.PluginContext, endpointCheckHealth)
 		res, err := a.checkHealthHandler.CheckHealth(ctx, parsedReq)
 		if err != nil {
 			return nil, err
diff --git a/backend/instancemgmt/instance_manager.go b/backend/instancemgmt/instance_manager.go
index 08e1d06..4ca761d 100644
--- a/backend/instancemgmt/instance_manager.go
+++ b/backend/instancemgmt/instance_manager.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"reflect"
 	"sync"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -17,6 +18,7 @@ var (
 		Name:      "active_instances",
 		Help:      "The number of active plugin instances",
 	})
+	disposeTTL = 5 * time.Second
 )
 
 // Instance is a marker interface for an instance.
@@ -119,7 +121,11 @@ func (im *instanceManager) Get(ctx context.Context, pluginContext backend.Plugin
 		}
 
 		if disposer, valid := ci.instance.(InstanceDisposer); valid {
-			disposer.Dispose()
+			time.AfterFunc(disposeTTL, func() {
+				disposer.Dispose()
+				activeInstances.Dec()
+			})
+		} else {
 			activeInstances.Dec()
 		}
 	}
diff --git a/backend/log.go b/backend/log.go
index 4e08f39..9e80b08 100644
--- a/backend/log.go
+++ b/backend/log.go
@@ -1,9 +1,23 @@
 package backend
 
 import (
+	"context"
+
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 )
 
+const (
+	endpointCallResource    = "callResource"
+	endpointCheckHealth     = "checkHealth"
+	endpointCollectMetrics  = "collectMetrics"
+	endpointQueryData       = "queryData"
+	endpointRunStream       = "runStream"
+	endpointSubscribeStream = "subscribeStream"
+	endpointPublishStream   = "publishStream"
+)
+
 // Logger is the default logger instance. This can be used directly to log from
 // your plugin to grafana-server with calls like backend.Logger.Debug(...).
 var Logger = log.DefaultLogger
@@ -12,3 +26,22 @@ var Logger = log.DefaultLogger
 var NewLoggerWith = func(args ...interface{}) log.Logger {
 	return log.New().With(args...)
 }
+
+func withContextualLogAttributes(ctx context.Context, pCtx PluginContext, endpoint string) context.Context {
+	args := []any{"pluginID", pCtx.PluginID, "endpoint", endpoint}
+	if tid := trace.SpanContextFromContext(ctx).TraceID(); tid.IsValid() {
+		args = append(args, "traceID", tid.String())
+	}
+	if pCtx.DataSourceInstanceSettings != nil {
+		args = append(
+			args,
+			"dsName", pCtx.DataSourceInstanceSettings.Name,
+			"dsUID", pCtx.DataSourceInstanceSettings.UID,
+		)
+		if pCtx.User != nil {
+			args = append(args, "uname", pCtx.User.Name)
+		}
+	}
+	ctx = log.WithContextualAttributes(ctx, args)
+	return ctx
+}
diff --git a/backend/log/context.go b/backend/log/context.go
new file mode 100644
index 0000000..87a6f44
--- /dev/null
+++ b/backend/log/context.go
@@ -0,0 +1,26 @@
+package log
+
+import "context"
+
+type loggerParamsCtxKeyType struct{}
+
+var loggerParamsCtxKey = loggerParamsCtxKeyType{}
+
+// WithContextualAttributes returns a new context with the given key/value log parameters appended to the existing ones.
+// It's possible to get a logger with those contextual parameters by using [FromContext].
+func WithContextualAttributes(ctx context.Context, logParams []any) context.Context {
+	p := logParams
+	if ctxParams := ctx.Value(loggerParamsCtxKey); ctxParams != nil {
+		p = append(ctxParams.([]any), logParams...)
+	}
+	return context.WithValue(ctx, loggerParamsCtxKey, p)
+}
+
+// ContextualAttributesFromContext returns the contextual key/value log parameters from the given context.
+// If no contextual log parameters are set, it returns nil.
+func ContextualAttributesFromContext(ctx context.Context) []any {
+	if logParams := ctx.Value(loggerParamsCtxKey); logParams != nil {
+		return logParams.([]any)
+	}
+	return nil
+}
diff --git a/backend/log/log.go b/backend/log/log.go
index e23185a..cf893da 100644
--- a/backend/log/log.go
+++ b/backend/log/log.go
@@ -2,6 +2,8 @@
 package log
 
 import (
+	"context"
+
 	hclog "github.com/hashicorp/go-hclog"
 )
 
@@ -24,6 +26,7 @@ type Logger interface {
 	Error(msg string, args ...interface{})
 	With(args ...interface{}) Logger
 	Level() Level
+	FromContext(ctx context.Context) Logger
 }
 
 // New creates a new logger.
@@ -90,5 +93,11 @@ func (l *hclogWrapper) With(args ...interface{}) Logger {
 	}
 }
 
+// FromContext creates a sub-logger with the contextual log parameters from the given context.
+// The contextual log parameters can be set using [WithContextualAttributes].
+func (l *hclogWrapper) FromContext(ctx context.Context) Logger {
+	return l.With(ContextualAttributesFromContext(ctx)...)
+}
+
 // DefaultLogger is the default logger.
 var DefaultLogger = New()
diff --git a/backend/resource_adapter.go b/backend/resource_adapter.go
index 7a56fa3..1f85e8c 100644
--- a/backend/resource_adapter.go
+++ b/backend/resource_adapter.go
@@ -39,5 +39,6 @@ func (a *resourceSDKAdapter) CallResource(protoReq *pluginv2.CallResourceRequest
 	ctx = WithGrafanaConfig(ctx, NewGrafanaCfg(protoReq.PluginContext.GrafanaConfig))
 	parsedReq := FromProto().CallResourceRequest(protoReq)
 	ctx = withHeaderMiddleware(ctx, parsedReq.GetHTTPHeaders())
+	ctx = withContextualLogAttributes(ctx, parsedReq.PluginContext, endpointCallResource)
 	return a.callResourceHandler.CallResource(ctx, parsedReq, fn)
 }
diff --git a/backend/stream_adapter.go b/backend/stream_adapter.go
index e985667..1e781fc 100644
--- a/backend/stream_adapter.go
+++ b/backend/stream_adapter.go
@@ -26,7 +26,9 @@ func (a *streamSDKAdapter) SubscribeStream(ctx context.Context, protoReq *plugin
 	}
 	ctx = propagateTenantIDIfPresent(ctx)
 	ctx = WithGrafanaConfig(ctx, NewGrafanaCfg(protoReq.PluginContext.GrafanaConfig))
-	resp, err := a.streamHandler.SubscribeStream(ctx, FromProto().SubscribeStreamRequest(protoReq))
+	parsedReq := FromProto().SubscribeStreamRequest(protoReq)
+	ctx = withContextualLogAttributes(ctx, parsedReq.PluginContext, endpointSubscribeStream)
+	resp, err := a.streamHandler.SubscribeStream(ctx, parsedReq)
 	if err != nil {
 		return nil, err
 	}
@@ -39,7 +41,9 @@ func (a *streamSDKAdapter) PublishStream(ctx context.Context, protoReq *pluginv2
 	}
 	ctx = propagateTenantIDIfPresent(ctx)
 	ctx = WithGrafanaConfig(ctx, NewGrafanaCfg(protoReq.PluginContext.GrafanaConfig))
-	resp, err := a.streamHandler.PublishStream(ctx, FromProto().PublishStreamRequest(protoReq))
+	parsedReq := FromProto().PublishStreamRequest(protoReq)
+	ctx = withContextualLogAttributes(ctx, parsedReq.PluginContext, endpointPublishStream)
+	resp, err := a.streamHandler.PublishStream(ctx, parsedReq)
 	if err != nil {
 		return nil, err
 	}
@@ -61,6 +65,8 @@ func (a *streamSDKAdapter) RunStream(protoReq *pluginv2.RunStreamRequest, protoS
 	ctx := protoSrv.Context()
 	ctx = propagateTenantIDIfPresent(ctx)
 	ctx = WithGrafanaConfig(ctx, NewGrafanaCfg(protoReq.PluginContext.GrafanaConfig))
+	parsedReq := FromProto().RunStreamRequest(protoReq)
+	ctx = withContextualLogAttributes(ctx, parsedReq.PluginContext, endpointRunStream)
 	sender := NewStreamSender(&runStreamServer{protoSrv: protoSrv})
-	return a.streamHandler.RunStream(ctx, FromProto().RunStreamRequest(protoReq), sender)
+	return a.streamHandler.RunStream(ctx, parsedReq, sender)
 }


Test patch:
diff --git a/backend/instancemgmt/instance_manager_test.go b/backend/instancemgmt/instance_manager_test.go
index e86e67c..8351fe2 100644
--- a/backend/instancemgmt/instance_manager_test.go
+++ b/backend/instancemgmt/instance_manager_test.go
@@ -43,6 +43,11 @@ func TestInstanceManager(t *testing.T) {
 					Updated: time.Now(),
 				},
 			}
+			origDisposeTTL := disposeTTL
+			disposeTTL = time.Millisecond
+			t.Cleanup(func() {
+				disposeTTL = origDisposeTTL
+			})
 			newInstance, err := im.Get(ctx, pCtxUpdated)
 
 			t.Run("New instance should be created", func(t *testing.T) {
@@ -57,7 +62,9 @@ func TestInstanceManager(t *testing.T) {
 			})
 
 			t.Run("Old instance should be disposed", func(t *testing.T) {
-				require.True(t, instance.(*testInstance).disposed)
+				instance.(*testInstance).wg.Wait()
+				require.True(t, instance.(*testInstance).disposed.Load())
+				require.Equal(t, int64(1), instance.(*testInstance).disposedTimes.Load())
 			})
 		})
 	})
@@ -95,7 +102,7 @@ func TestInstanceManagerConcurrency(t *testing.T) {
 		t.Run("All created instances should be either disposed or exist in cache for later disposing", func(t *testing.T) {
 			cachedInstance, _ := im.Get(ctx, pCtx)
 			for _, instance := range createdInstances {
-				if cachedInstance.(*testInstance) != instance && instance.disposedTimes < 1 {
+				if cachedInstance.(*testInstance) != instance && instance.disposedTimes.Load() < 1 {
 					require.FailNow(t, "Found lost reference to un-disposed instance")
 				}
 			}
@@ -103,6 +110,12 @@ func TestInstanceManagerConcurrency(t *testing.T) {
 	})
 
 	t.Run("Check possible race condition issues when re-creating instance on settings update", func(t *testing.T) {
+		origDisposeTTL := disposeTTL
+		disposeTTL = time.Millisecond
+		t.Cleanup(func() {
+			disposeTTL = origDisposeTTL
+		})
+
 		ctx := context.Background()
 		initialPCtx := backend.PluginContext{
 			OrgID: 1,
@@ -141,12 +154,13 @@ func TestInstanceManagerConcurrency(t *testing.T) {
 		wg.Wait()
 
 		t.Run("Initial instance should be disposed only once", func(t *testing.T) {
-			require.Equal(t, int64(1), instanceToDispose.(*testInstance).disposedTimes, "Instance should be disposed only once")
+			instanceToDispose.(*testInstance).wg.Wait()
+			require.Equal(t, int64(1), instanceToDispose.(*testInstance).disposedTimes.Load(), "Instance should be disposed only once")
 		})
 		t.Run("All created instances should be either disposed or exist in cache for later disposing", func(t *testing.T) {
 			cachedInstance, _ := im.Get(ctx, updatedPCtx)
 			for _, instance := range createdInstances {
-				if cachedInstance.(*testInstance) != instance && instance.disposedTimes < 1 {
+				if cachedInstance.(*testInstance) != instance && instance.disposedTimes.Load() < 1 {
 					require.FailNow(t, "Found lost reference to un-disposed instance")
 				}
 			}
@@ -204,13 +218,15 @@ func TestInstanceManagerConcurrency(t *testing.T) {
 type testInstance struct {
 	orgID         int64
 	updated       time.Time
-	disposed      bool
-	disposedTimes int64
+	disposed      atomic.Bool
+	disposedTimes atomic.Int64
+	wg            sync.WaitGroup
 }
 
 func (ti *testInstance) Dispose() {
-	ti.disposed = true
-	atomic.AddInt64(&ti.disposedTimes, 1)
+	ti.disposed.Store(true)
+	ti.disposedTimes.Add(1)
+	ti.wg.Done()
 }
 
 type testInstanceProvider struct {
@@ -231,8 +247,12 @@ func (tip *testInstanceProvider) NewInstance(_ context.Context, pluginContext ba
 	if tip.delay > 0 {
 		time.Sleep(tip.delay)
 	}
-	return &testInstance{
+
+	ti := &testInstance{
 		orgID:   pluginContext.OrgID,
 		updated: pluginContext.AppInstanceSettings.Updated,
-	}, nil
+	}
+	ti.wg.Add(1)
+
+	return ti, nil
 }
diff --git a/backend/log/context_test.go b/backend/log/context_test.go
new file mode 100644
index 0000000..3cd5246
--- /dev/null
+++ b/backend/log/context_test.go
@@ -0,0 +1,31 @@
+package log
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestContextualLogger(t *testing.T) {
+	t.Run("WithContextualAttributes", func(t *testing.T) {
+		t.Run("simple", func(t *testing.T) {
+			logParams := []any{"key", "value"}
+			ctx := WithContextualAttributes(context.Background(), logParams)
+			attrs := ContextualAttributesFromContext(ctx)
+			require.Equal(t, logParams, attrs)
+		})
+
+		t.Run("should append to existing value", func(t *testing.T) {
+			ctx := WithContextualAttributes(context.Background(), []any{"a", "b"})
+			ctx = WithContextualAttributes(ctx, []any{"c", "d"})
+			attrs := ContextualAttributesFromContext(ctx)
+			require.Equal(t, []any{"a", "b", "c", "d"}, attrs)
+		})
+	})
+
+	t.Run("ContextualAttributesFromContext on empty context should return empty slice", func(t *testing.T) {
+		attrs := ContextualAttributesFromContext(context.Background())
+		require.Empty(t, attrs)
+	})
+}
diff --git a/backend/log/log_test.go b/backend/log/log_test.go
index 7874893..cd157bf 100644
--- a/backend/log/log_test.go
+++ b/backend/log/log_test.go
@@ -1,20 +1,46 @@
-package log_test
+package log
 
 import (
+	"context"
 	"testing"
 
-	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestLogLevel(t *testing.T) {
-	logger := log.New()
+	logger := New()
 	level := logger.Level()
-	assert.Equal(t, level, log.Debug)
+	assert.Equal(t, level, Debug)
 }
 
 func TestLogLevelWarn(t *testing.T) {
-	logger := log.NewWithLevel(log.Warn)
+	logger := NewWithLevel(Warn)
 	level := logger.Level()
-	assert.Equal(t, level, log.Warn)
+	assert.Equal(t, level, Warn)
+}
+
+func TestLogFromContext(t *testing.T) {
+	t.Run("should return a new empty logger with empty context", func(t *testing.T) {
+		ctx := context.Background()
+		logger := New()
+		ctxLogger := logger.FromContext(ctx)
+		require.NotEqual(t, logger, ctxLogger)
+		require.Empty(t, ctxLogger.(*hclogWrapper).logger.ImpliedArgs())
+	})
+
+	t.Run("should return logger with contextual params when set in context", func(t *testing.T) {
+		ctx := WithContextualAttributes(context.Background(), []any{"key", "value"})
+		logger := New()
+		ctxLogger := logger.FromContext(ctx)
+		require.NotEqual(t, logger, ctxLogger)
+		require.Equal(t, []any{"key", "value"}, ctxLogger.(*hclogWrapper).logger.ImpliedArgs())
+	})
+
+	t.Run("params set in the logger should be kept in the contextual logger", func(t *testing.T) {
+		logger := New().With("service", "myservice")
+		ctxLogger := logger.FromContext(WithContextualAttributes(context.Background(), []any{"key", "value"}))
+		require.NotEqual(t, logger, ctxLogger)
+		require.Equal(t, []any{"key", "value", "service", "myservice"}, ctxLogger.(*hclogWrapper).logger.ImpliedArgs())
+	})
 }
diff --git a/backend/log_test.go b/backend/log_test.go
new file mode 100644
index 0000000..cb76e8a
--- /dev/null
+++ b/backend/log_test.go
@@ -0,0 +1,125 @@
+package backend
+
+import (
+	"context"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
+	"github.com/grafana/grafana-plugin-sdk-go/genproto/pluginv2"
+)
+
+func checkCtxLogger(ctx context.Context, t *testing.T, expParams map[string]any) {
+	t.Helper()
+	logAttrs := log.ContextualAttributesFromContext(ctx)
+	if len(expParams) == 0 {
+		require.Empty(t, logAttrs)
+		return
+	}
+
+	require.NotEmpty(t, logAttrs)
+	require.Truef(t, len(logAttrs)%2 == 0, "expected even number of log params, got %d", len(logAttrs))
+	require.Equal(t, len(expParams)*2, len(logAttrs))
+	for i := 0; i < len(logAttrs)/2; i++ {
+		key := logAttrs[i*2].(string)
+		val := logAttrs[i*2+1]
+		expVal, ok := expParams[key]
+		require.Truef(t, ok, "unexpected log param: %s", key)
+		require.Equal(t, expVal, val)
+	}
+}
+
+func TestContextualLogger(t *testing.T) {
+	const pluginID = "plugin-id"
+	pCtx := &pluginv2.PluginContext{PluginId: pluginID}
+	t.Run("DataSDKAdapter", func(t *testing.T) {
+		run := make(chan struct{}, 1)
+		a := newDataSDKAdapter(QueryDataHandlerFunc(func(ctx context.Context, req *QueryDataRequest) (*QueryDataResponse, error) {
+			checkCtxLogger(ctx, t, map[string]any{"endpoint": "queryData", "pluginID": pluginID})
+			run <- struct{}{}
+			return NewQueryDataResponse(), nil
+		}))
+		_, err := a.QueryData(context.Background(), &pluginv2.QueryDataRequest{
+			PluginContext: pCtx,
+		})
+		require.NoError(t, err)
+		<-run
+	})
+
+	t.Run("DiagnosticsSDKAdapter", func(t *testing.T) {
+		run := make(chan struct{}, 1)
+		a := newDiagnosticsSDKAdapter(prometheus.DefaultGatherer, CheckHealthHandlerFunc(func(ctx context.Context, req *CheckHealthRequest) (*CheckHealthResult, error) {
+			checkCtxLogger(ctx, t, map[string]any{"endpoint": "checkHealth", "pluginID": pluginID})
+			run <- struct{}{}
+			return &CheckHealthResult{}, nil
+		}))
+		_, err := a.CheckHealth(context.Background(), &pluginv2.CheckHealthRequest{
+			PluginContext: pCtx,
+		})
+		require.NoError(t, err)
+		<-run
+	})
+
+	t.Run("ResourceSDKAdapter", func(t *testing.T) {
+		run := make(chan struct{}, 1)
+		a := newResourceSDKAdapter(CallResourceHandlerFunc(func(ctx context.Context, req *CallResourceRequest, sender CallResourceResponseSender) error {
+			checkCtxLogger(ctx, t, map[string]any{"endpoint": "callResource", "pluginID": pluginID})
+			run <- struct{}{}
+			return nil
+		}))
+		err := a.CallResource(&pluginv2.CallResourceRequest{
+			PluginContext: pCtx,
+		}, newTestCallResourceServer())
+		require.NoError(t, err)
+		<-run
+	})
+
+	t.Run("StreamHandler", func(t *testing.T) {
+		subscribeStreamRun := make(chan struct{}, 1)
+		publishStreamRun := make(chan struct{}, 1)
+		runStreamRun := make(chan struct{}, 1)
+		a := newStreamSDKAdapter(&streamAdapter{
+			subscribeStreamFunc: func(ctx context.Context, request *SubscribeStreamRequest) (*SubscribeStreamResponse, error) {
+				checkCtxLogger(ctx, t, map[string]any{"endpoint": "subscribeStream", "pluginID": pluginID})
+				subscribeStreamRun <- struct{}{}
+				return &SubscribeStreamResponse{}, nil
+			},
+			publishStreamFunc: func(ctx context.Context, request *PublishStreamRequest) (*PublishStreamResponse, error) {
+				checkCtxLogger(ctx, t, map[string]any{"endpoint": "publishStream", "pluginID": pluginID})
+				publishStreamRun <- struct{}{}
+				return &PublishStreamResponse{}, nil
+			},
+			runStreamFunc: func(ctx context.Context, request *RunStreamRequest, sender *StreamSender) error {
+				checkCtxLogger(ctx, t, map[string]any{"endpoint": "runStream", "pluginID": pluginID})
+				runStreamRun <- struct{}{}
+				return nil
+			},
+		})
+
+		t.Run("SubscribeStream", func(t *testing.T) {
+			_, err := a.SubscribeStream(context.Background(), &pluginv2.SubscribeStreamRequest{
+				PluginContext: pCtx,
+			})
+			require.NoError(t, err)
+			<-subscribeStreamRun
+		})
+
+		t.Run("PublishStream", func(t *testing.T) {
+			_, err := a.PublishStream(context.Background(), &pluginv2.PublishStreamRequest{
+				PluginContext: pCtx,
+			})
+			require.NoError(t, err)
+			<-publishStreamRun
+		})
+
+		t.Run("RunStream", func(t *testing.T) {
+			err := a.RunStream(&pluginv2.RunStreamRequest{
+				PluginContext: pCtx,
+			}, newTestRunStreamServer())
+			require.NoError(t, err)
+			<-runStreamRun
+		})
+	})
+}
diff --git a/experimental/e2e/storage/har_test.go b/experimental/e2e/storage/har_test.go
index 84d6088..cc8245a 100644
--- a/experimental/e2e/storage/har_test.go
+++ b/experimental/e2e/storage/har_test.go
@@ -44,11 +44,16 @@ func TestHARStorage(t *testing.T) {
 				require.NoError(t, e)
 				c <- true
 			}()
-			req, res := exampleRequest()
-			defer res.Body.Close()
-			req.URL.Path = "/two"
-			err = two.Add(req, res)
-			require.NoError(t, err)
+			go func() {
+				req, res := exampleRequest()
+				defer res.Body.Close()
+				req.URL.Path = "/two"
+				err = two.Add(req, res)
+				require.NoError(t, err)
+				c <- true
+			}()
+
+			<-c
 			<-c
 			require.Len(t, one.Entries(), 2)
 			require.Len(t, two.Entries(), 2)