collector: bound MI probe queries in Build with a timeout (#2657)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/collector/dns/dns.go

@@ -322,7 +322,7 @@ func (c *Collector) buildErrorStatsCollector(miSession *mi.Session) error {
 	c.miQuery = query
 
 	var stats []Statistic
-	if err := c.miSession.Query(&stats, mi.NamespaceRootMicrosoftDNS, c.miQuery, 0); err != nil {
+	if err := c.miSession.Query(&stats, mi.NamespaceRootMicrosoftDNS, c.miQuery, mi.BuildQueryTimeout); err != nil {
 		return fmt.Errorf("failed to query DNS statistics: %w", err)
 	}
 

internal/collector/fsrmquota/fsrmquota.go

@@ -94,7 +94,7 @@ func (c *Collector) Build(_ *slog.Logger, miSession *mi.Session) error {
 	c.buildDescriptors()
 
 	var dst []msftFSRMQuota
-	if err := c.miSession.Query(&dst, mi.NamespaceRootWindowsFSRM, c.miQuery, 0); err != nil {
+	if err := c.miSession.Query(&dst, mi.NamespaceRootWindowsFSRM, c.miQuery, mi.BuildQueryTimeout); err != nil {
 		return fmt.Errorf("WMI query failed: %w", err)
 	}
 

internal/collector/process/process.go

@@ -322,7 +322,7 @@ func (c *Collector) buildWorkerProcessWMI(miSession *mi.Session) error {
 
 	var workerProcesses []WorkerProcess
 
-	if err = miSession.Query(&workerProcesses, mi.NamespaceRootWebAdministration, miQuery, 0); err != nil {
+	if err = miSession.Query(&workerProcesses, mi.NamespaceRootWebAdministration, miQuery, mi.BuildQueryTimeout); err != nil {
 		c.logger.LogAttrs(context.Background(), slog.LevelDebug,
 			`root\WebAdministration isn't available, reading IIS application pool names from the w3wp command line only`,
 			slog.Any("err", err),

internal/collector/storage_spaces/storage_spaces_pool.go

@@ -107,7 +107,7 @@ func (c *Collector) buildStoragePool() error {
 
 	var dst []msftStoragePool
 
-	if err := c.miSession.Query(&dst, mi.NamespaceRootStorage, c.storagePoolMIQuery, 0); err != nil {
+	if err := c.miSession.Query(&dst, mi.NamespaceRootStorage, c.storagePoolMIQuery, mi.BuildQueryTimeout); err != nil {
 		return fmt.Errorf("WMI query failed: %w", err)
 	}
 

internal/collector/storage_spaces/storage_spaces_virtual_disk.go

@@ -104,7 +104,7 @@ func (c *Collector) buildVirtualDisk() error {
 
 	var dst []msftVirtualDisk
 
-	if err := c.miSession.Query(&dst, mi.NamespaceRootStorage, c.virtualDiskMIQuery, 0); err != nil {
+	if err := c.miSession.Query(&dst, mi.NamespaceRootStorage, c.virtualDiskMIQuery, mi.BuildQueryTimeout); err != nil {
 		return fmt.Errorf("WMI query failed: %w", err)
 	}
 

internal/mi/mi_test.go

@@ -772,6 +772,35 @@ func Test_MI_QueryTimeout(t *testing.T) {
 	require.NoError(t, err)
 }
 
+// Test_MI_Session_QueryTimeout checks that the timeout passed to Session.Query,
+// e.g. mi.BuildQueryTimeout, reaches MI and surfaces as the native result.
+func Test_MI_Session_QueryTimeout(t *testing.T) {
+	application, err := mi.ApplicationInitialize()
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, application.Close()) })
+
+	destinationOptions, err := application.NewDestinationOptions()
+	require.NoError(t, err)
+	require.NoError(t, destinationOptions.SetLocale(mi.LocaleEnglish))
+
+	session, err := application.NewSession(destinationOptions)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, session.Close()) })
+
+	query, err := mi.NewQuery("select * from win32_process")
+	require.NoError(t, err)
+
+	var processes []win32Process
+
+	err = session.Query(&processes, mi.NamespaceRootCIMv2, query, time.Millisecond)
+	require.ErrorIs(t, err, mi.MI_RESULT_INVALID_OPERATION_TIMEOUT)
+
+	processes = nil
+
+	require.NoError(t, session.Query(&processes, mi.NamespaceRootCIMv2, query, mi.BuildQueryTimeout))
+	require.NotEmpty(t, processes)
+}
+
 func Test_MI_Query_Uint16Array(t *testing.T) {
 	application, err := mi.ApplicationInitialize()
 	require.NoError(t, err)

internal/mi/session.go

@@ -29,6 +29,12 @@ import (
 	"golang.org/x/sys/windows"
 )
 
+// BuildQueryTimeout bounds the probe queries collectors run in Build. Build
+// has no scrape budget, and a zero timeout sets no MI timeout, so a stalled
+// provider would otherwise block startup. A timed-out query fails with
+// MI_RESULT_INVALID_OPERATION_TIMEOUT. MI ignores timeouts below 1ms.
+const BuildQueryTimeout = 30 * time.Second
+
 // Session represents a session.
 //
 // https://learn.microsoft.com/en-us/windows/win32/api/mi/ns-mi-mi_session