From 20612bb81943118fae581b06fc268f8207bdc6ec Mon Sep 17 00:00:00 2001
From: Ashutosh Narkar <anarkar4387@gmail.com>
Date: Thu, 6 Jul 2023 18:23:52 -0700
Subject: [PATCH] rego: Parse store modules iff modules set on the Rego object

Currently we parse store modules irrespective of whether
there are modules on the Rego object. This will result in
the compilation of those modules which triggers the bundle
activation flow. Now as part of the module compilation
we interact with the compiler's modules and run compilation
on the input modules. If let's say there are concurrent health
check requests (ie. /v1/health), this could result in a race
during the compilation process while working with the compiler's
modules.

This change avoids this situation by skipping parsing of the store
modules when none are set on the Rego object. The assumption this
change makes is that while using the rego package the compiler and
store are kept in-sync.

Fixes: #5868

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>
---
 rego/rego.go          | 12 +++++-------
 sdk/opa_test.go       |  8 ++++----
 server/server_test.go |  1 -
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/rego/rego.go b/rego/rego.go
index 73b34c28ed..0e17667a83 100644
--- a/rego/rego.go
+++ b/rego/rego.go
@@ -1736,22 +1736,20 @@ func (r *Rego) prepare(ctx context.Context, qType queryType, extras []extraStage
 }
 
 func (r *Rego) parseModules(ctx context.Context, txn storage.Transaction, m metrics.Metrics) error {
+	if len(r.modules) == 0 {
+		return nil
+	}
+
 	ids, err := r.store.ListPolicies(ctx, txn)
 	if err != nil {
 		return err
 	}
 
-	// if there are no raw modules, nor modules in the store, then there
-	// is nothing to do.
-	if len(r.modules) == 0 && len(ids) == 0 {
-		return nil
-	}
-
 	m.Timer(metrics.RegoModuleParse).Start()
 	defer m.Timer(metrics.RegoModuleParse).Stop()
 	var errs Errors
 
-	// Parse any modules in the are saved to the store, but only if
+	// Parse any modules that are saved to the store, but only if
 	// another compile step is going to occur (ie. we have parsed modules
 	// that need to be compiled).
 	for _, id := range ids {
diff --git a/sdk/opa_test.go b/sdk/opa_test.go
index 63e89bc76b..0551fc2abe 100644
--- a/sdk/opa_test.go
+++ b/sdk/opa_test.go
@@ -500,7 +500,7 @@ main = true
 		t.Fatal("expected true but got:", decision, ok)
 	}
 
-	if exp, act := 5, len(m.All()); exp != act {
+	if exp, act := 4, len(m.All()); exp != act {
 		t.Fatalf("expected %d metrics, got %d", exp, act)
 	}
 
@@ -576,7 +576,7 @@ main = true
 		t.Fatal("expected true but got:", decision, ok)
 	}
 
-	if exp, act := 26, len(m.All()); exp != act {
+	if exp, act := 25, len(m.All()); exp != act {
 		t.Fatalf("expected %d metrics, got %d", exp, act)
 	}
 
@@ -1024,7 +1024,7 @@ allow {
 		t.Fatal("expected &{[2 = data.junk.x] []} true but got:", decision, ok)
 	}
 
-	if exp, act := 6, len(m.All()); exp != act {
+	if exp, act := 5, len(m.All()); exp != act {
 		t.Fatalf("expected %d metrics, got %d", exp, act)
 	}
 
@@ -1111,7 +1111,7 @@ allow {
 		t.Fatal("expected &{[2 = data.junk.x] []} true but got:", decision, ok)
 	}
 
-	if exp, act := 33, len(m.All()); exp != act {
+	if exp, act := 32, len(m.All()); exp != act {
 		t.Fatalf("expected %d metrics, got %d", exp, act)
 	}
 
diff --git a/server/server_test.go b/server/server_test.go
index 3335b2601a..3f69f6a9ee 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -1038,7 +1038,6 @@ func TestCompileV1Observability(t *testing.T) {
 			"timer_rego_partial_eval_ns",
 			"timer_rego_query_compile_ns",
 			"timer_rego_query_parse_ns",
-			"timer_rego_module_parse_ns",
 			"timer_server_handler_ns",
 			"counter_disk_read_keys",
 			"timer_disk_read_ns",