Allow OPA compiler state to be saved and restored

[anderseknert]

One awesome aspect of Go is its excellent faciliities for embedding resources along with compiled executables. We piggyback on this capabilty in Regal to embed the entire bundle directory of the projects, containing ~300 Rego policies and a few JSON/YAML data files. This number will of course only grow with time. For pretty much all commands provided, our ideal is to drop out of Go as quickly as possible and do as much processing as we possibly can in the wonderful world of Rego ✨ But before we can do that, we must first compile the embedded bundle. This (compiling the Regal bundle) takes ~100 ms on my beefy M4 Max laptop, and of course considerably more on a more modestly specced machine, and dramatically more on under-specced CI runners.

Now to be clear: a few hundred milliseconds is nothing in terms of compilation time compared to many other languages, and this isn't generally a bottleneck for the typical OPA use case. However, for applications like Regal that essentially are customized OPA + embedded policies, what becomes a real problem is the fact that compilation needs to be done every time the executable runs. It doesn't matter that once we ship a binary, the embedded content is unchanged!

While this may be a non-issue for a long-running process like a server, even a few hundred milliseconds of delay before any actual processing takes place is a big deal for CLI applications running one-off commands, as they'll be perceived as sluggish. Now that the Rego evaluation is incredibly fast, this is quite a shame. So let's fix that if we can!

One obvious thing to try is of course to improve the performance of the compiler, and there's much to explore in that space (like e.g. concurrency #8060). For the embedded use case specifically, I'd however like to suggest that we also explore the option of allowing compiled state to be persisted and later loaded into the compiler in order to avoid doing much work at all. Essentially — we should allow advanced users and use cases a way to say "this is a bundle you have already compiled, load it without going through any compiler stages that aren't already in the compiled output". The most costly aspect of compilation is the many stages of rewriting Rego, and performing these stages should be entirely avoidable when provided policies where this has already been done. Some state — like the mapping of variable names to their rewritten aliases — is AFAIK not part of the compiler's output today. But we could easily include additional files containing the metadata necessary to quickly restore compiler state as part of the output from compilation when a "persist mode" toggle is set. Perhaps it could all fit in a bundle .manifest, or we'd add new files for this purpose. I'm sure we'll figure it out.

In my always humble opinion, this would make embedding Rego a truly amazing option for building OPA-powered CLI apps 😃

[srenatus]

ℹ️ For CLI applications that evaluate Rego with policies that don't change during the live time of a specific version of the CLI app, using Wasm might be an interesting option. Evaluating a wasm-bundle via v1/rego doesn't require the compiler to do much, so it can be thought of as "compiled policy code". It should be feasible to bake a wasm bundle into a CLI app.

[anderseknert]

I can imagine that is a good alternative in some cases, but for various reasons not a viable solution for Regal.

Just out of curiosity — what would be required for "Evaluating a wasm-bundle via v1/rego" in an application integrating OPA as a dependency? I don't see this use case covered in our docs on Wasm. I assume I'd have to provide the runtime (Wasmtime only?) myself, and build my application with that embedded. Sounds doable, but quite inconvenient. Or I can use Wazero for eval? If that works it'd be a great addition to our docs 😃

[srenatus]

Yeah this isn't easy to discover without searching through the code. The gist is:

1. go build -tags=opa_wasm to get wasmtime-go pull into OPA
2. opa build -t wasm ..., go:embed the bundle
3. pass the bundle path to v1/rego, rego.LoadBundle("path/to/bundle.tar.gz")

From then on, plain rego.New(...).Eval() should do the trick, transparently evaluate the wasm bundle using wasmtime-go.

Now I'm curious if it really works 😅

[srenatus]

Update: It works, and this is how. I had forgotten about the feature import in my message above. Also, we don't need the filesystem in the binary if we load the bundle ourselves.

main.go:

package main

import (
	"bytes"
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"log"

	"github.com/open-policy-agent/opa/v1/bundle"
	_ "github.com/open-policy-agent/opa/v1/features/wasm"
	"github.com/open-policy-agent/opa/v1/rego"
)

//go:embed bundle.tar.gz
var bundleBytes []byte

func main() {
	ctx := context.Background()

	reader := bytes.NewReader(bundleBytes)
	bundleReader := bundle.NewCustomReader(bundle.NewTarballLoader(reader))

	b, err := bundleReader.Read()
	if err != nil {
		log.Fatal(err)
	}

	input := map[string]any{
		"user": "admin",
	}

	r := rego.New(
		rego.Query("data.example"),
		rego.ParsedBundle("mybundle", &b),
		rego.Input(input),
	)

	rs, err := r.Eval(ctx)
	if err != nil {
		log.Fatal(err)
	}

	result, _ := json.MarshalIndent(rs, "", "  ")
	fmt.Printf("%s\n", result)
}

policy.rego:

package example
default allow := false
allow if input.user == "admin"

1. run opa build -t wasm policy.rego, this creates the bundle.tar.gz
2. build: go build -tags opa_wasm -o wasm-demo main.go
3. run: ./wasm-demo:

$ ./wasm-demo
[
  {
    "expressions": [
      {
        "value": {
          "allow": true
        },
        "text": "data.example",
        "location": {
          "row": 1,
          "col": 1
        }
      }
    ]
  }
]

[anderseknert]

Awesome! We should add this to the OPA docs on Wasm 😃