From b4baf3a30922f7f82dcad8282d055410ac759f94 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan-Otto=20Kr=C3=B6pke?= <mail@jkroepke.de>
Date: Sun, 21 Apr 2024 17:44:31 +0200
Subject: [PATCH] scheduled_task: Move OLE connection to collector
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>
---
 .../scheduled_task/scheduled_task.go          | 31 +++++++++++--------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/pkg/collector/scheduled_task/scheduled_task.go b/pkg/collector/scheduled_task/scheduled_task.go
index 047f85500..e6d77a695 100644
--- a/pkg/collector/scheduled_task/scheduled_task.go
+++ b/pkg/collector/scheduled_task/scheduled_task.go
@@ -3,6 +3,7 @@
 package scheduled_task
 
 import (
+	"errors"
 	"fmt"
 	"regexp"
 	"runtime"
@@ -63,7 +64,6 @@ const (
 	TASK_RESULT_SUCCESS TaskResult = 0x0
 )
 
-// RegisteredTask ...
 type ScheduledTask struct {
 	Name            string
 	Path            string
@@ -117,18 +117,6 @@ func (c *collector) GetPerfCounter() ([]string, error) {
 }
 
 func (c *collector) Build() error {
-	runtime.LockOSThread()
-	defer runtime.UnlockOSThread()
-
-	err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED)
-	if err != nil {
-		code := err.(*ole.OleError).Code()
-		if code != ole.S_OK && code != S_FALSE {
-			return err
-		}
-	}
-	defer ole.CoUninitialize()
-
 	c.LastResult = prometheus.NewDesc(
 		prometheus.BuildFQName(types.Namespace, Name, "last_result"),
 		"The result that was returned the last time the registered task was run",
@@ -150,6 +138,8 @@ func (c *collector) Build() error {
 		nil,
 	)
 
+	var err error
+
 	c.taskIncludePattern, err = regexp.Compile(fmt.Sprintf("^(?:%s)$", *c.taskInclude))
 	if err != nil {
 		return err
@@ -231,6 +221,21 @@ const SCHEDULED_TASK_PROGRAM_ID = "Schedule.Service.1"
 const S_FALSE = 0x00000001
 
 func getScheduledTasks() (scheduledTasks ScheduledTasks, err error) {
+	// The only way to run WMI queries in parallel while being thread-safe is to
+	// ensure the CoInitialize[Ex]() call is bound to its current OS thread.
+	// Otherwise, attempting to initialize and run parallel queries across
+	// goroutines will result in protected memory errors.
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+
+	if err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED); err != nil {
+		var oleCode *ole.OleError
+		if errors.As(err, &oleCode) && oleCode.Code() != ole.S_OK && oleCode.Code() != S_FALSE {
+			return nil, err
+		}
+	}
+	defer ole.CoUninitialize()
+
 	schedClassID, err := ole.ClassIDFrom(SCHEDULED_TASK_PROGRAM_ID)
 	if err != nil {
 		return scheduledTasks, err