testing: Improve lsp shutdown logic & address test race conditions (#1943)

* Improve server shutdown logic address test race conditions

I have been trying to make test suite flakes less common. I am not 100% sure I have got all the issues, but these changes make the flakes and races appear much less often on my local testing. The main things done:

- Add proper shutdown with timeout and worker synchronization. Refactor server test helper to manage the shutdown of the server instances centrally. This helps address issues where workers were continuing to write to pipes after closing.
- Address Aggregate Data races (ast.Object sharing), uses defensive copies when setting aggregate data. ast.Object.Insert() contains non-atomic map/slice operations.
- Address channel blocking issues and timeouts, e.g. in TestTemplateWorkerRaceCondition
- Config access races, locking when accessing workspaceRootURI.
- Fixer state mutation race, create local copy of OPAFmtOpts before mutation.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Fix race condition on client field access

Add clientLock mutex to protect LanguageServer.client from concurrent
access. The jsonrpc2 asyncHandler's separate goroutines for initialize
and initialized requests caused a race when workers read client info
after initializationGate closes.

Follows same pattern as loadedConfigLock.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Replace workspace lint limit with named constant

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Rename JobRouter to LintJobRouter

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Load aggregates from data.workspace in ls

Enable aggregate rules to work both in normal linting mode (aggregates
via input) and when aggregates are provided via shared store
(data.workspace.aggregates).

This unifies aggregate loading in CLI and LSP by moving aggregate
storage from cache to the lsp regoStore.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Refactor LSP init and parallelize file loading

Most importantly, we now do a full workspace lint to avoid linting files
one by one, once this is done, the single file job worker starts.

Replace OverwriteAggregates/AggregateReportOnly flags with clearer
IsInitialization/FullWorkspaceLint to fix init ordering and reduce races.
File lint workers now wait for the initial workspace lint before processing,
and workspace loading runs across 10 concurrent goroutines.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Remove StartHoverWorker call from test setup

Missed after rebase updates from #1951

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Remove old waitForDiagnostics helper

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Address test flakes in lsp server tests

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Address data race with locking in bundle cache

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Reduce test polling time for server tests

This makes some of the slower polling tests faster as we learn about
state changes quicker.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Update ls to use a one-at-a-time linter runner

This has been done to make the server more maintainable.

We have found that the partial linting with aggregates to be complex and hard
to maintain. At the same time, the linter is a lot faster now so the impact is
minimal.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Inline lint job runner into diagnostics worker

Remove the lintJobRunner abstraction and move its logic directly into
StartDiagnosticsWorker, simplifying shutdown and reducing indirection.

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

* Remove extra workspace lint job

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

---------

Signed-off-by: Charlie Egan <charlie_egan@apple.com>

Author: charlieegan3

bundle/regal/main/main.rego

@@ -233,7 +233,7 @@ _remove_empty_aggregates(aggregates) := {"aggregate": set()} if {
 #   - input: schema.regal.aggregate
 _aggregate_report_inputs[cat_title].aggregate contains formatted if {
 	some filename, cat_title
-	aggregate := input.aggregates_internal[filename][cat_title][_]
+	aggregate := aggregates_internal[filename][cat_title][_]
 	formatted := _format_aggregate(aggregate, filename)
 }
 
@@ -252,3 +252,12 @@ _ignored(violation, directives) if {
 	ignored_rules := directives[violation.location.row + 1]
 	violation.title in ignored_rules
 }
+
+# METADATA
+# description: |
+#   returns aggregates_internal from input if available, otherwise falls back to data.workspace.aggregates
+#   this allows aggregate rules to work both in normal linting mode and when aggregates are provided via data
+# scope: rule
+aggregates_internal := input.aggregates_internal
+
+aggregates_internal := data.workspace.aggregates if not input.aggregates_internal

bundle/regal/rules/bugs/impossible-not/impossible_not.rego

@@ -6,6 +6,7 @@
 package regal.rules.bugs["impossible-not"]
 
 import data.regal.ast
+import data.regal.main
 import data.regal.result
 import data.regal.util
 
@@ -144,5 +145,5 @@ _resolve(ref, pkg_path, imported_symbols) := concat(".", resolved) if {
 _aggregates[file] := agg if {
 	# only one aggregate exported per file, so this is safe
 	some file
-	agg := input.aggregates_internal[file]["bugs/impossible-not"][_]
+	agg := main.aggregates_internal[file]["bugs/impossible-not"][_]
 }

bundle/regal/rules/custom/missing-metadata/missing_metadata.rego

@@ -7,6 +7,7 @@ package regal.rules.custom["missing-metadata"]
 
 import data.regal.ast
 import data.regal.config
+import data.regal.main
 import data.regal.result
 import data.regal.util
 
@@ -97,7 +98,7 @@ aggregate_report contains violation if {
 _package_path_aggs[pkg_name] contains [file, item] if {
 	some file, items in _aggregates
 	some item in items
-	pkg_name := input.aggregates_internal[file].common.package_name
+	pkg_name := main.aggregates_internal[file].common.package_name
 }
 
 # METADATA
@@ -122,4 +123,4 @@ _excepted_rule_pattern(cfg, value) if regex.match(cfg["except-rule-path-pattern"
 # METADATA
 # schemas:
 #   - input: schema.regal.aggregate
-_aggregates[file] := input.aggregates_internal[file]["custom/missing-metadata"] if some file
+_aggregates[file] := main.aggregates_internal[file]["custom/missing-metadata"] if some file

bundle/regal/rules/idiomatic/no-defined-entrypoint/no_defined_entrypoint.rego

@@ -5,6 +5,7 @@
 #     ref: https://www.openpolicyagent.org/projects/regal/rules/idiomatic/no-defined-entrypoint
 package regal.rules.idiomatic["no-defined-entrypoint"]
 
+import data.regal.main
 import data.regal.result
 import data.regal.util
 
@@ -38,4 +39,4 @@ aggregate_report contains violation if {
 # METADATA
 # schemas:
 #   - input: schema.regal.aggregate
-_any_entrypoint if input.aggregates_internal[_]["idiomatic/no-defined-entrypoint"][_].entrypoint
+_any_entrypoint if main.aggregates_internal[_]["idiomatic/no-defined-entrypoint"][_].entrypoint

bundle/regal/rules/idiomatic/no-defined-entrypoint/no_defined_entrypoint_test.rego

@@ -64,3 +64,34 @@ test_success_multiple_entrypoints_defined if {
 
 	r == set()
 }
+
+# Test the fallback pattern where aggregates come from data.workspace.aggregates
+# instead of input.aggregates_internal - expects success (no violations)
+test_success_entrypoint_from_data_workspace_aggregates if {
+	# Using with data.workspace.aggregates instead of with input.aggregates_internal
+	r := rule.aggregate_report with data.workspace.aggregates as {"p.rego": {"idiomatic/no-defined-entrypoint": [
+		{"entrypoint": {"col": 1, "file": "policy.rego", "row": 2}},
+		{"entrypoint": {"col": 1, "file": "policy.rego", "row": 6}},
+	]}}
+
+	r == set()
+}
+
+# Test the fallback actually works - expects failure (violation) when no entrypoint is defined
+test_fail_no_entrypoint_with_data_workspace if {
+	# Explicitly ensure input.aggregates_internal is undefined by using an empty input
+	# No entrypoint defined, so should fail
+	r := rule.aggregate_report with input as {}
+		with data.workspace.aggregates as {"p.rego": {"idiomatic/no-defined-entrypoint": []}}
+
+	r == {{
+		"category": "idiomatic",
+		"description": "Missing entrypoint annotation",
+		"level": "error",
+		"related_resources": [{
+			"description": "documentation",
+			"ref": "https://www.openpolicyagent.org/projects/regal/rules/idiomatic/no-defined-entrypoint",
+		}],
+		"title": "no-defined-entrypoint",
+	}}
+}

bundle/regal/rules/imports/circular-import/circular_import.rego

@@ -8,6 +8,7 @@
 package regal.rules.imports["circular-import"]
 
 import data.regal.ast
+import data.regal.main
 import data.regal.result
 import data.regal.util
 
@@ -117,5 +118,5 @@ _aggregated[file] := agg if {
 	# we know that there is only one aggregate of this type per file,
 	# so we can simplify things some for our callers
 	some file
-	agg := input.aggregates_internal[file]["imports/circular-import"][_]
+	agg := main.aggregates_internal[file]["imports/circular-import"][_]
 }

bundle/regal/rules/imports/prefer-package-imports/prefer_package_imports.rego

@@ -7,6 +7,7 @@ package regal.rules.imports["prefer-package-imports"]
 
 import data.regal.ast
 import data.regal.config
+import data.regal.main
 import data.regal.result
 import data.regal.util
 
@@ -77,5 +78,5 @@ _aggregates[file] := agg if {
 	# we know that there is only one aggregate of this type per file,
 	# so we can simplify things some for our callers
 	some file
-	agg := input.aggregates_internal[file]["imports/prefer-package-imports"][_]
+	agg := main.aggregates_internal[file]["imports/prefer-package-imports"][_]
 }

bundle/regal/rules/imports/unresolved-import/unresolved_import.rego

@@ -7,6 +7,7 @@ package regal.rules.imports["unresolved-import"]
 
 import data.regal.ast
 import data.regal.config
+import data.regal.main
 import data.regal.result
 import data.regal.util
 
@@ -46,10 +47,10 @@ aggregate contains entry if {
 # schemas:
 #   - input: schema.regal.aggregate
 aggregate_report contains violation if {
-	all_known_refs := {path | path := input.aggregates_internal[_]["imports/unresolved-import"][_].exported_refs[_]}
+	all_known_refs := {path | path := main.aggregates_internal[_]["imports/unresolved-import"][_].exported_refs[_]}
 
 	some file
-	entry := input.aggregates_internal[file]["imports/unresolved-import"][_]
+	entry := main.aggregates_internal[file]["imports/unresolved-import"][_]
 
 	some [path, location] in entry.imports
 	not path in _except_imports
@@ -108,5 +109,5 @@ _aggregates[file] := agg if {
 	# we know that there is only one aggregate of this type per file,
 	# so we can simplify things some for our callers
 	some file
-	agg := input.aggregates_internal[file]["imports/unresolved-import"][_]
+	agg := main.aggregates_internal[file]["imports/unresolved-import"][_]
 }

bundle/regal/rules/imports/unresolved-reference/unresolved_reference.rego

@@ -7,6 +7,7 @@ package regal.rules.imports["unresolved-reference"]
 
 import data.regal.ast
 import data.regal.config
+import data.regal.main
 import data.regal.result
 
 # METADATA
@@ -153,4 +154,4 @@ _to_location_object(loc, text, file) := {"location": {
 # METADATA
 # schemas:
 #   - input: schema.regal.aggregate
-_aggregates[file] := input.aggregates_internal[file]["imports/unresolved-reference"] if some file
+_aggregates[file] := main.aggregates_internal[file]["imports/unresolved-reference"] if some file

cmd/languageserver.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/spf13/cobra"
 
@@ -27,7 +28,6 @@ func init() {
 
 		RunE: wrapProfiling(func([]string) error {
 			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
 
 			if exe, err := os.Executable(); err != nil {
 				fmt.Fprintln(os.Stderr, "error getting executable:", err)
@@ -62,14 +62,14 @@ func init() {
 
 			ls.SetConn(conn)
 
-			go ls.StartDiagnosticsWorker(ctx)
-			go ls.StartTestLocationsWorker(ctx)
-			go ls.StartCommandWorker(ctx)
-			go ls.StartConfigWorker(ctx)
-			go ls.StartWorkspaceStateWorker(ctx)
-			go ls.StartTemplateWorker(ctx)
-			go ls.StartQueryCacheWorker(ctx)
-			go ls.StartWebServer(ctx)
+			ls.StartDiagnosticsWorker(ctx)
+			ls.StartTestLocationsWorker(ctx)
+			ls.StartCommandWorker(ctx)
+			ls.StartConfigWorker(ctx)
+			ls.StartWorkspaceStateWorker(ctx)
+			ls.StartTemplateWorker(ctx)
+			ls.StartQueryCacheWorker(ctx)
+			ls.StartWebServer(ctx)
 
 			sigChan := make(chan os.Signal, 1)
 			signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
@@ -81,6 +81,15 @@ func init() {
 				fmt.Fprintln(os.Stderr, "signal: ", sig.String())
 			}
 
+			cancel()
+
+			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer shutdownCancel()
+
+			if err := ls.Shutdown(shutdownCtx); err != nil {
+				fmt.Fprintln(os.Stderr, "shutdown error:", err)
+			}
+
 			return nil
 		}),
 	}