fn: agent call overrider (#1080)

Similar to LB Agent call overrider, this PR adds Agent overrider
for Agents to modify/analyze a Call/Extensions during GetCall().

api/agent/agent.go

@@ -120,6 +120,8 @@ type agent struct {
 	shutonce            sync.Once
 	callEndCount        int64
 	disableAsyncDequeue bool
+
+	callOverrider CallOverrider
 }
 
 type AgentOption func(*agent) error
@@ -194,6 +196,17 @@ func WithoutAsyncDequeue() AgentOption {
 	}
 }
 
+// Agents can use this to register a CallOverrider to modify a Call and extensions
+func WithCallOverrider(fn CallOverrider) AgentOption {
+	return func(a *agent) error {
+		if a.callOverrider != nil {
+			return errors.New("lb-agent call overriders already exists")
+		}
+		a.callOverrider = fn
+		return nil
+	}
+}
+
 // Create a default docker driver from agent config
 func NewDockerDriver(cfg *AgentConfig) *docker.DockerDriver {
 	return docker.NewDocker(drivers.Config{

api/agent/call.go

@@ -41,6 +41,9 @@ type Call interface {
 	End(ctx context.Context, err error) error
 }
 
+// Interceptor in GetCall
+type CallOverrider func(*models.Call, map[string]string) (map[string]string, error)
+
 // TODO build w/o closures... lazy
 type CallOpt func(c *call) error
 
@@ -261,6 +264,16 @@ func (a *agent) GetCall(opts ...CallOpt) (Call, error) {
 		return nil, errors.New("no model or request provided for call")
 	}
 
+	// If overrider is present, let's allow it to modify models.Call
+	// and call extensions
+	if a.callOverrider != nil {
+		ext, err := a.callOverrider(c.Call, c.extensions)
+		if err != nil {
+			return nil, err
+		}
+		c.extensions = ext
+	}
+
 	mem := c.Memory + uint64(c.TmpFsSize)
 	if !a.resources.IsResourcePossible(mem, uint64(c.CPUs), c.Type == models.TypeAsync) {
 		// if we're not going to be able to run this call on this machine, bail here.

api/agent/lb_agent.go

@@ -17,8 +17,6 @@ import (
 	"github.com/fnproject/fn/fnext"
 )
 
-type CallOverrider func(*models.Call, map[string]string) (map[string]string, error)
-
 type lbAgent struct {
 	cfg           AgentConfig
 	da            DataAccess
@@ -40,7 +38,7 @@ func WithLBAgentConfig(cfg *AgentConfig) LBAgentOption {
 }
 
 // LB agents can use this to register a CallOverrider to modify a Call and extensions
-func WithCallOverrider(fn CallOverrider) LBAgentOption {
+func WithLBCallOverrider(fn CallOverrider) LBAgentOption {
 	return func(a *lbAgent) error {
 		if a.callOverrider != nil {
 			return errors.New("lb-agent call overriders already exists")

test/fn-system-tests/exec_test.go

@@ -110,6 +110,12 @@ func TestCanExecuteFunction(t *testing.T) {
 	if err != nil || cheese != "Tete de Moine" {
 		t.Fatalf("getConfigContent/FN_CHEESE check failed (%v) on %v", err, output)
 	}
+
+	// Now let's check FN_WINE, since runners have override to insert this.
+	wine, err := getConfigContent("FN_WINE", output.Bytes())
+	if err != nil || wine != "1982 Margaux" {
+		t.Fatalf("getConfigContent/FN_WINE check failed (%v) on %v", err, output)
+	}
 }
 
 func TestCanExecuteBigOutput(t *testing.T) {

test/fn-system-tests/system_test.go

@@ -226,7 +226,7 @@ func SetUpLBNode(ctx context.Context) (*server.Server, error) {
 
 	// Create an LB Agent with a Call Overrider to intercept calls in GetCall(). Overrider in this example
 	// scrubs CPU/TmpFsSize and adds FN_CHEESE key/value into extensions.
-	lbAgent, err := agent.NewLBAgent(agent.NewCachedDataAccess(cl), nodePool, placer, agent.WithCallOverrider(LBCallOverrider))
+	lbAgent, err := agent.NewLBAgent(agent.NewCachedDataAccess(cl), nodePool, placer, agent.WithLBCallOverrider(LBCallOverrider))
 	if err != nil {
 		return nil, err
 	}
@@ -266,7 +266,11 @@ func SetUpPureRunnerNode(ctx context.Context, nodeNum int) (*server.Server, erro
 	}
 
 	// inner agent for pure-runners
-	innerAgent := agent.New(ds, agent.WithConfig(cfg), agent.WithDockerDriver(drv), agent.WithoutAsyncDequeue())
+	innerAgent := agent.New(ds,
+		agent.WithConfig(cfg),
+		agent.WithDockerDriver(drv),
+		agent.WithoutAsyncDequeue(),
+		agent.WithCallOverrider(PureRunnerCallOverrider))
 
 	cancelCtx, cancel := context.WithCancel(ctx)
 
@@ -370,6 +374,18 @@ func LBCallOverrider(c *models.Call, exts map[string]string) (map[string]string,
 	return exts, nil
 }
 
+// Pure Runner Agent Call Option
+func PureRunnerCallOverrider(c *models.Call, exts map[string]string) (map[string]string, error) {
+
+	if exts == nil {
+		exts = make(map[string]string)
+	}
+
+	// Add an FN_WINE extension, just an example...
+	exts["FN_WINE"] = "1982 Margaux"
+	return exts, nil
+}
+
 // An example Pure Runner docker driver. Using CreateCookie, it intercepts a generated cookie to
 // add an environment variable FN_CHEESE if it finds a FN_CHEESE extension.
 type customDriver struct {
@@ -383,20 +399,26 @@ func (d *customDriver) CreateCookie(ctx context.Context, task drivers.ContainerT
 		return cookie, err
 	}
 
-	// if call extensions include 'foo', then let's add FN_CHEESE env vars, which should
-	// end up in Env/Config.
-	ext := task.Extensions()
-	cheese, ok := ext["FN_CHEESE"]
-	if ok {
 	// docker driver specific data
 	obj := cookie.ContainerOptions()
 	opts, ok := obj.(docker.CreateContainerOptions)
 	if !ok {
 		logrus.Fatal("Unexpected driver, should be docker")
 	}
+
+	// if call extensions include 'foo', then let's add FN_CHEESE env vars, which should
+	// end up in Env/Config.
+	ext := task.Extensions()
+	cheese, ok := ext["FN_CHEESE"]
+	if ok {
 		opts.Config.Env = append(opts.Config.Env, "FN_CHEESE="+cheese)
 	}
 
+	wine, ok := ext["FN_WINE"]
+	if ok {
+		opts.Config.Env = append(opts.Config.Env, "FN_WINE="+wine)
+	}
+
 	return cookie, nil
 }