fix(jsonnet): add configurable timeout (#798)

ory · Jul 19, 2024 · 287b4d9 · 287b4d9
1 parent 471f7d6
commit 287b4d9
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 19 deletions.
diff --git a/jsonnetsecure/jsonnet.go b/jsonnetsecure/jsonnet.go
@@ -9,11 +9,18 @@ import (
 	"path"
 	"runtime"
 	"testing"
+	"time"
 
 	"github.com/google/go-jsonnet"
 )
 
 type (
+	evaluationOptions struct {
+		evalTimeout time.Duration
+	}
+
+	EvaluationOptionModifier func(*evaluationOptions)
+
 	VM interface {
 		EvaluateAnonymousSnippet(filename string, snippet string) (json string, formattedErr error)
 		ExtCode(key string, val string)
@@ -31,10 +38,11 @@ type (
 	}
 
 	ProcessVM struct {
-		ctx    context.Context
-		path   string
-		args   []string
-		params processParameters
+		ctx         context.Context
+		path        string
+		args        []string
+		execTimeout time.Duration
+		params      processParameters
 	}
 
 	vmOptions struct {
@@ -43,6 +51,7 @@ type (
 		args              []string
 		ctx               context.Context
 		pool              *pool
+		execTimeout       time.Duration
 	}
 
 	Option func(o *vmOptions)
@@ -70,6 +79,12 @@ func WithProcessIsolatedVM(ctx context.Context) Option {
 	}
 }
 
+func WithExecTimeout(timeout time.Duration) Option {
+	return func(o *vmOptions) {
+		o.execTimeout = timeout
+	}
+}
+
 func WithJsonnetBinary(jsonnetBinaryPath string) Option {
 	return func(o *vmOptions) {
 		o.jsonnetBinaryPath = jsonnetBinaryPath

diff --git a/jsonnetsecure/jsonnet_pool.go b/jsonnetsecure/jsonnet_pool.go
@@ -5,6 +5,7 @@ package jsonnetsecure
 
 import (
 	"bufio"
+	"cmp"
 	"context"
 	"encoding/json"
 	"io"
@@ -23,11 +24,12 @@ import (
 
 type (
 	processPoolVM struct {
-		path   string
-		args   []string
-		ctx    context.Context
-		params processParameters
-		pool   *pool
+		path        string
+		args        []string
+		ctx         context.Context
+		params      processParameters
+		execTimeout time.Duration
+		pool        *pool
 	}
 	Pool interface {
 		Close()
@@ -183,7 +185,7 @@ func (vm *processPoolVM) EvaluateAnonymousSnippet(filename string, snippet strin
 	defer otelx.End(span, &err)
 
 	// TODO: maybe leave the timeout to the caller?
-	ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, cmp.Or(vm.execTimeout, 1*time.Second))
 	defer cancel()
 
 	params := vm.params
@@ -222,10 +224,11 @@ func NewProcessPoolVM(opts *vmOptions) VM {
 		ctx = context.Background()
 	}
 	return &processPoolVM{
-		path: opts.jsonnetBinaryPath,
-		args: opts.args,
-		ctx:  ctx,
-		pool: opts.pool,
+		path:        opts.jsonnetBinaryPath,
+		args:        opts.args,
+		ctx:         ctx,
+		pool:        opts.pool,
+		execTimeout: opts.execTimeout,
 	}
 }
 

diff --git a/jsonnetsecure/jsonnet_processvm.go b/jsonnetsecure/jsonnet_processvm.go
@@ -5,6 +5,7 @@ package jsonnetsecure
 
 import (
 	"bytes"
+	"cmp"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -23,9 +24,10 @@ import (
 
 func NewProcessVM(opts *vmOptions) VM {
 	return &ProcessVM{
-		path: opts.jsonnetBinaryPath,
-		args: opts.args,
-		ctx:  opts.ctx,
+		path:        opts.jsonnetBinaryPath,
+		args:        opts.args,
+		ctx:         opts.ctx,
+		execTimeout: opts.execTimeout,
 	}
 }
 
@@ -35,12 +37,11 @@ func (p *ProcessVM) EvaluateAnonymousSnippet(filename string, snippet string) (_
 	defer otelx.End(span, &err)
 
 	// We retry the process creation, because it sometimes times out.
-	const processVMTimeout = 1 * time.Second
 	return backoff.RetryWithData(func() (_ string, err error) {
 		ctx, span := tracer.Start(ctx, "jsonnetsecure.ProcessVM.EvaluateAnonymousSnippet.run")
 		defer otelx.End(span, &err)
 
-		ctx, cancel := context.WithTimeout(ctx, processVMTimeout)
+		ctx, cancel := context.WithTimeout(ctx, cmp.Or(p.execTimeout, 1*time.Second))
 		defer cancel()
 
 		var (