From 337998d58b256dce6ce70be8af08641ec8edc70e Mon Sep 17 00:00:00 2001
From: Marcus Efraimsson <marcus.efraimsson@gmail.com>
Date: Thu, 15 Jul 2021 14:33:38 +0200
Subject: [PATCH] Chore: Ease the migration of always using context.Context
 when interacting with the bus (#36733)

If DispatchCtx is called and there's a handler registered using
AddHandler we now allow that and not return a handler not
found error. In addition we log a warning message helping
the developers know that AddHandlerCtx should be used
instead.
If Dispatch is called and there's a handler registered
using AddHandlerCtx we log a warning message helping
the developers know that DispatchCtx should be used instead.
---
 pkg/bus/bus.go      | 23 ++++++++++++++++++++---
 pkg/bus/bus_test.go | 16 ++++++++++++++++
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/pkg/bus/bus.go b/pkg/bus/bus.go
index 7958b955f31..6dd8a745604 100644
--- a/pkg/bus/bus.go
+++ b/pkg/bus/bus.go
@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"reflect"
 
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/setting"
 	"github.com/opentracing/opentracing-go"
 )
 
@@ -52,6 +54,7 @@ func (b *InProcBus) InTransaction(ctx context.Context, fn func(ctx context.Conte
 
 // InProcBus defines the bus structure
 type InProcBus struct {
+	logger          log.Logger
 	handlers        map[string]HandlerFunc
 	handlersWithCtx map[string]HandlerFunc
 	listeners       map[string][]HandlerFunc
@@ -63,7 +66,9 @@ var globalBus = New()
 
 // New initialize the bus
 func New() Bus {
-	bus := &InProcBus{}
+	bus := &InProcBus{
+		logger: log.New("bus"),
+	}
 	bus.handlers = make(map[string]HandlerFunc)
 	bus.handlersWithCtx = make(map[string]HandlerFunc)
 	bus.listeners = make(map[string][]HandlerFunc)
@@ -91,13 +96,22 @@ func (b *InProcBus) DispatchCtx(ctx context.Context, msg Msg) error {
 
 	span.SetTag("msg", msgName)
 
+	withCtx := true
 	var handler = b.handlersWithCtx[msgName]
 	if handler == nil {
-		return ErrHandlerNotFound
+		withCtx = false
+		handler = b.handlers[msgName]
+		if handler == nil {
+			return ErrHandlerNotFound
+		}
 	}
 
 	var params = []reflect.Value{}
-	params = append(params, reflect.ValueOf(ctx))
+	if withCtx {
+		params = append(params, reflect.ValueOf(ctx))
+	} else if setting.Env == setting.Dev {
+		b.logger.Warn("DispatchCtx called with message handler registered using AddHandler and should be changed to use AddHandlerCtx", "msgName", msgName)
+	}
 	params = append(params, reflect.ValueOf(msg))
 
 	ret := reflect.ValueOf(handler).Call(params)
@@ -124,6 +138,9 @@ func (b *InProcBus) Dispatch(msg Msg) error {
 
 	var params = []reflect.Value{}
 	if withCtx {
+		if setting.Env == setting.Dev {
+			b.logger.Warn("Dispatch called with message handler registered using AddHandlerCtx and should be changed to use DispatchCtx", "msgName", msgName)
+		}
 		params = append(params, reflect.ValueOf(context.Background()))
 	}
 	params = append(params, reflect.ValueOf(msg))
diff --git a/pkg/bus/bus_test.go b/pkg/bus/bus_test.go
index a0febee3223..38840e27798 100644
--- a/pkg/bus/bus_test.go
+++ b/pkg/bus/bus_test.go
@@ -69,6 +69,22 @@ func TestDispatchCtx(t *testing.T) {
 	require.True(t, invoked, "expected handler to be called")
 }
 
+func TestDispatchCtx_NoContextHandler(t *testing.T) {
+	bus := New()
+
+	var invoked bool
+
+	bus.AddHandler(func(query *testQuery) error {
+		invoked = true
+		return nil
+	})
+
+	err := bus.DispatchCtx(context.Background(), &testQuery{})
+	require.NoError(t, err)
+
+	require.True(t, invoked, "expected handler to be called")
+}
+
 func TestDispatchCtx_NoRegisteredHandler(t *testing.T) {
 	bus := New()