From 56aace8ae9248e89a632c370ff7526303b944152 Mon Sep 17 00:00:00 2001
From: Arne Luenser <arne.luenser@ory.sh>
Date: Tue, 30 Jul 2024 14:41:34 +0200
Subject: [PATCH] fix: jsonnet pool timeouts

The fixed 1 second timeout now only applies to the Jsonnet evaluation step,
but not the the spawning of the process which can occasionally take longer.

Additionally, we now warm the pool asynchronously on construction.
---
 jsonnetsecure/jsonnet_pool.go | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/jsonnetsecure/jsonnet_pool.go b/jsonnetsecure/jsonnet_pool.go
index ec395113..41ebcb2a 100644
--- a/jsonnetsecure/jsonnet_pool.go
+++ b/jsonnetsecure/jsonnet_pool.go
@@ -65,6 +65,10 @@ func NewProcessPool(size int) Pool {
 	if err != nil {
 		panic(err) // this should never happen, see implementation of puddle.NewPool
 	}
+	for range size {
+		// warm pool
+		go pud.CreateResource(context.Background())
+	}
 	go func() {
 		for {
 			time.Sleep(10 * time.Second)
@@ -141,12 +145,19 @@ func newWorker(ctx context.Context) (_ worker, err error) {
 	errs := make(chan string)
 	go scan(errs, stderr)
 
-	return worker{
+	w := worker{
 		cmd:    cmd,
 		stdin:  in,
 		stdout: out,
 		stderr: errs,
-	}, nil
+	}
+	_, err = w.eval(ctx, []byte("{}")) // warm up
+	if err != nil {
+		w.destroy()
+		return worker{}, errors.Wrap(err, "newWorker: warm up failed")
+	}
+
+	return w, nil
 }
 
 func (w worker) destroy() {
@@ -182,9 +193,6 @@ func (vm *processPoolVM) EvaluateAnonymousSnippet(filename string, snippet strin
 	ctx, span := tracer.Start(vm.ctx, "jsonnetsecure.processPoolVM.EvaluateAnonymousSnippet", trace.WithAttributes(attribute.String("filename", filename)))
 	defer otelx.End(span, &err)
 
-	// TODO: maybe leave the timeout to the caller?
-	ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
-	defer cancel()
 
 	params := vm.params
 	params.Filename = filename
@@ -201,6 +209,8 @@ func (vm *processPoolVM) EvaluateAnonymousSnippet(filename string, snippet strin
 		return "", errors.Wrap(err, "jsonnetsecure: acquire")
 	}
 
+	ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
+	defer cancel()
 	result, err := worker.Value().eval(ctx, pp)
 	if err != nil {
 		worker.Destroy()