add SERVICE_PATTERNS env var (#1153)

agent/systemd.go

@@ -27,6 +27,7 @@ type systemdManager struct {
 	serviceStatsMap map[string]*systemd.Service
 	isRunning       bool
 	hasFreshStats   bool
+	patterns        []string
 }
 
 // newSystemdManager creates a new systemdManager.
@@ -39,6 +40,7 @@ func newSystemdManager() (*systemdManager, error) {
 
 	manager := &systemdManager{
 		serviceStatsMap: make(map[string]*systemd.Service),
+		patterns:        getServicePatterns(),
 	}
 
 	manager.startWorker(conn)
@@ -109,7 +111,7 @@ func (sm *systemdManager) getServiceStats(conn *dbus.Conn, refresh bool) []*syst
 		defer conn.Close()
 	}
 
-	units, err := conn.ListUnitsByPatternsContext(context.Background(), []string{"loaded"}, []string{"*.service"})
+	units, err := conn.ListUnitsByPatternsContext(context.Background(), []string{"loaded"}, sm.patterns)
 	if err != nil {
 		slog.Error("Error listing systemd service units", "err", err)
 		return nil
@@ -245,3 +247,26 @@ func unescapeServiceName(name string) string {
 	}
 	return unescaped
 }
+
+// getServicePatterns returns the list of service patterns to match.
+// It reads from the SERVICE_PATTERNS environment variable if set,
+// otherwise defaults to "*service".
+func getServicePatterns() []string {
+	patterns := []string{}
+	if envPatterns, _ := GetEnv("SERVICE_PATTERNS"); envPatterns != "" {
+		for pattern := range strings.SplitSeq(envPatterns, ",") {
+			pattern = strings.TrimSpace(pattern)
+			if pattern == "" {
+				continue
+			}
+			if !strings.HasSuffix(pattern, ".service") {
+				pattern += ".service"
+			}
+			patterns = append(patterns, pattern)
+		}
+	}
+	if len(patterns) == 0 {
+		patterns = []string{"*.service"}
+	}
+	return patterns
+}

agent/systemd_test.go

@@ -3,6 +3,7 @@
 package agent
 
 import (
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -46,3 +47,112 @@ func TestUnescapeServiceNameInvalid(t *testing.T) {
 		})
 	}
 }
+
+func TestGetServicePatterns(t *testing.T) {
+	tests := []struct {
+		name           string
+		prefixedEnv    string
+		unprefixedEnv  string
+		expected       []string
+		cleanupEnvVars bool
+	}{
+		{
+			name:           "default when no env var set",
+			prefixedEnv:    "",
+			unprefixedEnv:  "",
+			expected:       []string{"*.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "single pattern with prefixed env",
+			prefixedEnv:    "nginx",
+			unprefixedEnv:  "",
+			expected:       []string{"nginx.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "single pattern with unprefixed env",
+			prefixedEnv:    "",
+			unprefixedEnv:  "nginx",
+			expected:       []string{"nginx.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "prefixed env takes precedence",
+			prefixedEnv:    "nginx",
+			unprefixedEnv:  "apache",
+			expected:       []string{"nginx.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "multiple patterns",
+			prefixedEnv:    "nginx,apache,postgresql",
+			unprefixedEnv:  "",
+			expected:       []string{"nginx.service", "apache.service", "postgresql.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "patterns with .service suffix",
+			prefixedEnv:    "nginx.service,apache.service",
+			unprefixedEnv:  "",
+			expected:       []string{"nginx.service", "apache.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "mixed patterns with and without suffix",
+			prefixedEnv:    "nginx.service,apache,postgresql.service",
+			unprefixedEnv:  "",
+			expected:       []string{"nginx.service", "apache.service", "postgresql.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "patterns with whitespace",
+			prefixedEnv:    " nginx , apache , postgresql ",
+			unprefixedEnv:  "",
+			expected:       []string{"nginx.service", "apache.service", "postgresql.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "empty patterns are skipped",
+			prefixedEnv:    "nginx,,apache,  ,postgresql",
+			unprefixedEnv:  "",
+			expected:       []string{"nginx.service", "apache.service", "postgresql.service"},
+			cleanupEnvVars: true,
+		},
+		{
+			name:           "wildcard pattern",
+			prefixedEnv:    "*nginx*,*apache*",
+			unprefixedEnv:  "",
+			expected:       []string{"*nginx*.service", "*apache*.service"},
+			cleanupEnvVars: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Clean up any existing env vars
+			os.Unsetenv("BESZEL_AGENT_SERVICE_PATTERNS")
+			os.Unsetenv("SERVICE_PATTERNS")
+
+			// Set up environment variables
+			if tt.prefixedEnv != "" {
+				os.Setenv("BESZEL_AGENT_SERVICE_PATTERNS", tt.prefixedEnv)
+			}
+			if tt.unprefixedEnv != "" {
+				os.Setenv("SERVICE_PATTERNS", tt.unprefixedEnv)
+			}
+
+			// Run the function
+			result := getServicePatterns()
+
+			// Verify results
+			assert.Equal(t, tt.expected, result, "Patterns should match expected values")
+
+			// Cleanup
+			if tt.cleanupEnvVars {
+				os.Unsetenv("BESZEL_AGENT_SERVICE_PATTERNS")
+				os.Unsetenv("SERVICE_PATTERNS")
+			}
+		})
+	}
+}