use ctx on routingFunc

app_test.go

@@ -285,19 +285,19 @@ func TestAddRoute(t *testing.T) {
 	initApp()
 	Configure(true, "testtype", Cluster, map[string]string{}, viper.New())
 	app.router = nil
-	err := AddRoute("somesv", func(session *session.Session, route *route.Route, payload []byte, servers map[string]*cluster.Server) (*cluster.Server, error) {
+	err := AddRoute("somesv", func(ctx context.Context, route *route.Route, payload []byte, servers map[string]*cluster.Server) (*cluster.Server, error) {
 		return nil, nil
 	})
 	assert.EqualError(t, constants.ErrRouterNotInitialized, err.Error())
 
 	app.router = router.New()
-	err = AddRoute("somesv", func(session *session.Session, route *route.Route, payload []byte, servers map[string]*cluster.Server) (*cluster.Server, error) {
+	err = AddRoute("somesv", func(ctx context.Context, route *route.Route, payload []byte, servers map[string]*cluster.Server) (*cluster.Server, error) {
 		return nil, nil
 	})
 	assert.NoError(t, err)
 
 	app.running = true
-	err = AddRoute("somesv", func(session *session.Session, route *route.Route, payload []byte, servers map[string]*cluster.Server) (*cluster.Server, error) {
+	err = AddRoute("somesv", func(ctx context.Context, route *route.Route, payload []byte, servers map[string]*cluster.Server) (*cluster.Server, error) {
 		return nil, nil
 	})
 	assert.EqualError(t, constants.ErrChangeRouteWhileRunning, err.Error())

router/router.go

@@ -21,6 +21,7 @@
 package router
 
 import (
+	"context"
 	"math/rand"
 	"time"
 
@@ -30,7 +31,6 @@ import (
 	"github.com/topfreegames/pitaya/logger"
 	"github.com/topfreegames/pitaya/protos"
 	"github.com/topfreegames/pitaya/route"
-	"github.com/topfreegames/pitaya/session"
 )
 
 // Router struct
@@ -41,7 +41,7 @@ type Router struct {
 
 // RoutingFunc defines a routing function
 type RoutingFunc func(
-	session *session.Session,
+	ctx context.Context,
 	route *route.Route,
 	payload []byte,
 	servers map[string]*cluster.Server,
@@ -74,9 +74,9 @@ func (r *Router) defaultRoute(
 
 // Route gets the right server to use in the call
 func (r *Router) Route(
+	ctx context.Context,
 	rpcType protos.RPCType,
 	svType string,
-	session *session.Session,
 	route *route.Route,
 	msg *message.Message,
 ) (*cluster.Server, error) {
@@ -97,7 +97,7 @@ func (r *Router) Route(
 		server := r.defaultRoute(serversOfType)
 		return server, nil
 	}
-	return routeFunc(session, route, msg.Data, serversOfType)
+	return routeFunc(ctx, route, msg.Data, serversOfType)
 }
 
 // AddRoute adds a routing function to a server type

router/router_test.go

@@ -1,6 +1,7 @@
 package router
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -11,7 +12,6 @@ import (
 	"github.com/topfreegames/pitaya/internal/message"
 	"github.com/topfreegames/pitaya/protos"
 	"github.com/topfreegames/pitaya/route"
-	"github.com/topfreegames/pitaya/session"
 )
 
 var (
@@ -24,7 +24,7 @@ var (
 	}
 
 	routingFunction = func(
-		session *session.Session,
+		ctx context.Context,
 		route *route.Route,
 		payload []byte,
 		servers map[string]*cluster.Server,
@@ -70,7 +70,7 @@ func TestDefaultRoute(t *testing.T) {
 func TestRoute(t *testing.T) {
 	t.Parallel()
 
-	session := &session.Session{}
+	ctx := context.Background()
 	route := route.NewRoute(serverType, "service", "method")
 
 	for name, table := range routerTables {
@@ -86,7 +86,7 @@ func TestRoute(t *testing.T) {
 			router.AddRoute(serverType, routingFunction)
 			router.SetServiceDiscovery(mockServiceDiscovery)
 
-			retServer, err := router.Route(table.rpcType, table.serverType, session, route, &message.Message{
+			retServer, err := router.Route(ctx, table.rpcType, table.serverType, route, &message.Message{
 				Data: []byte{0x01},
 			})
 			assert.Equal(t, table.server, retServer)

service/remote.go

@@ -426,7 +426,7 @@ func (r *RemoteService) remoteCall(
 	target := server
 
 	if target == nil {
-		target, err = r.router.Route(rpcType, svType, session, route, msg)
+		target, err = r.router.Route(ctx, rpcType, svType, route, msg)
 		if err != nil {
 			return nil, e.NewError(err, e.ErrInternalCode)
 		}