From 881a0ba1dbc9b0d2cc7dc651e0a3558331b3691e Mon Sep 17 00:00:00 2001 From: Tolga Ceylan Date: Wed, 20 Jun 2018 16:21:09 -0700 Subject: [PATCH] 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 | 13 ++++++++++ api/agent/call.go | 13 ++++++++++ api/agent/lb_agent.go | 4 +-- test/fn-system-tests/exec_test.go | 6 +++++ test/fn-system-tests/system_test.go | 38 +++++++++++++++++++++++------ 5 files changed, 63 insertions(+), 11 deletions(-) diff --git a/api/agent/agent.go b/api/agent/agent.go index 25b489552..ec9111418 100644 --- a/api/agent/agent.go +++ b/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{ diff --git a/api/agent/call.go b/api/agent/call.go index c387c3f00..b6a783e32 100644 --- a/api/agent/call.go +++ b/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. diff --git a/api/agent/lb_agent.go b/api/agent/lb_agent.go index dbea7fe9a..e8653a6c2 100644 --- a/api/agent/lb_agent.go +++ b/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") diff --git a/test/fn-system-tests/exec_test.go b/test/fn-system-tests/exec_test.go index 5acbb0bdd..8a015a768 100644 --- a/test/fn-system-tests/exec_test.go +++ b/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) { diff --git a/test/fn-system-tests/system_test.go b/test/fn-system-tests/system_test.go index d186d12e2..c2ddfe50b 100644 --- a/test/fn-system-tests/system_test.go +++ b/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 } + // 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 { - // docker driver specific data - obj := cookie.ContainerOptions() - opts, ok := obj.(docker.CreateContainerOptions) - if !ok { - logrus.Fatal("Unexpected driver, should be docker") - } 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 }