Fix memory leak by updating go-ole and adding wbem initialization (#82)

exporter.go

@@ -13,6 +13,7 @@ import (
 
 	"golang.org/x/sys/windows/svc"
 
+	"github.com/StackExchange/wmi"
 	"github.com/martinlindhe/wmi_exporter/collector"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -130,6 +131,18 @@ func init() {
 	prometheus.MustRegister(version.NewCollector("wmi_exporter"))
 }
 
+func initWbem() {
+	// This initialization prevents a memory leak on WMF 5+. See
+	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
+	// for details.
+	log.Debugf("Initializing SWbemServices")
+	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
+	if err != nil {
+		log.Fatal(err)
+	}
+	wmi.DefaultClient.SWbemServicesClient = s
+}
+
 func main() {
 	var (
 		showVersion       = flag.Bool("version", false, "Print version information.")
@@ -158,6 +171,8 @@ func main() {
 		return
 	}
 
+	initWbem()
+
 	isInteractive, err := svc.IsAnInteractiveSession()
 	if err != nil {
 		log.Fatal(err)

vendor/github.com/StackExchange/wmi/README.md

@@ -1,4 +1,6 @@
 wmi
 ===
 
-Package wmi provides a WQL interface for WMI on Windows.
+Package wmi provides a WQL interface to Windows WMI.
+
+Note: It interfaces with WMI on the local machine, therefore it only runs on Windows.

vendor/github.com/StackExchange/wmi/swbemservices.go

@@ -0,0 +1,260 @@
+// +build windows
+
+package wmi
+
+import (
+	"fmt"
+	"reflect"
+	"runtime"
+	"sync"
+
+	"github.com/go-ole/go-ole"
+	"github.com/go-ole/go-ole/oleutil"
+)
+
+// SWbemServices is used to access wmi. See https://msdn.microsoft.com/en-us/library/aa393719(v=vs.85).aspx
+type SWbemServices struct {
+	//TODO: track namespace. Not sure if we can re connect to a different namespace using the same instance
+	cWMIClient            *Client //This could also be an embedded struct, but then we would need to branch on Client vs SWbemServices in the Query method
+	sWbemLocatorIUnknown  *ole.IUnknown
+	sWbemLocatorIDispatch *ole.IDispatch
+	queries               chan *queryRequest
+	closeError            chan error
+	lQueryorClose         sync.Mutex
+}
+
+type queryRequest struct {
+	query    string
+	dst      interface{}
+	args     []interface{}
+	finished chan error
+}
+
+// InitializeSWbemServices will return a new SWbemServices object that can be used to query WMI
+func InitializeSWbemServices(c *Client, connectServerArgs ...interface{}) (*SWbemServices, error) {
+	//fmt.Println("InitializeSWbemServices: Starting")
+	//TODO: implement connectServerArgs as optional argument for init with connectServer call
+	s := new(SWbemServices)
+	s.cWMIClient = c
+	s.queries = make(chan *queryRequest)
+	initError := make(chan error)
+	go s.process(initError)
+
+	err, ok := <-initError
+	if ok {
+		return nil, err //Send error to caller
+	}
+	//fmt.Println("InitializeSWbemServices: Finished")
+	return s, nil
+}
+
+// Close will clear and release all of the SWbemServices resources
+func (s *SWbemServices) Close() error {
+	s.lQueryorClose.Lock()
+	if s == nil || s.sWbemLocatorIDispatch == nil {
+		s.lQueryorClose.Unlock()
+		return fmt.Errorf("SWbemServices is not Initialized")
+	}
+	if s.queries == nil {
+		s.lQueryorClose.Unlock()
+		return fmt.Errorf("SWbemServices has been closed")
+	}
+	//fmt.Println("Close: sending close request")
+	var result error
+	ce := make(chan error)
+	s.closeError = ce //Race condition if multiple callers to close. May need to lock here
+	close(s.queries)  //Tell background to shut things down
+	s.lQueryorClose.Unlock()
+	err, ok := <-ce
+	if ok {
+		result = err
+	}
+	//fmt.Println("Close: finished")
+	return result
+}
+
+func (s *SWbemServices) process(initError chan error) {
+	//fmt.Println("process: starting background thread initialization")
+	//All OLE/WMI calls must happen on the same initialized thead, so lock this goroutine
+	runtime.LockOSThread()
+	defer runtime.LockOSThread()
+
+	err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED)
+	if err != nil {
+		oleCode := err.(*ole.OleError).Code()
+		if oleCode != ole.S_OK && oleCode != S_FALSE {
+			initError <- fmt.Errorf("ole.CoInitializeEx error: %v", err)
+			return
+		}
+	}
+	defer ole.CoUninitialize()
+
+	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
+	if err != nil {
+		initError <- fmt.Errorf("CreateObject SWbemLocator error: %v", err)
+		return
+	} else if unknown == nil {
+		initError <- ErrNilCreateObject
+		return
+	}
+	defer unknown.Release()
+	s.sWbemLocatorIUnknown = unknown
+
+	dispatch, err := s.sWbemLocatorIUnknown.QueryInterface(ole.IID_IDispatch)
+	if err != nil {
+		initError <- fmt.Errorf("SWbemLocator QueryInterface error: %v", err)
+		return
+	}
+	defer dispatch.Release()
+	s.sWbemLocatorIDispatch = dispatch
+
+	// we can't do the ConnectServer call outside the loop unless we find a way to track and re-init the connectServerArgs
+	//fmt.Println("process: initialized. closing initError")
+	close(initError)
+	//fmt.Println("process: waiting for queries")
+	for q := range s.queries {
+		//fmt.Printf("process: new query: len(query)=%d\n", len(q.query))
+		errQuery := s.queryBackground(q)
+		//fmt.Println("process: s.queryBackground finished")
+		if errQuery != nil {
+			q.finished <- errQuery
+		}
+		close(q.finished)
+	}
+	//fmt.Println("process: queries channel closed")
+	s.queries = nil //set channel to nil so we know it is closed
+	//TODO: I think the Release/Clear calls can panic if things are in a bad state.
+	//TODO: May need to recover from panics and send error to method caller instead.
+	close(s.closeError)
+}
+
+// Query runs the WQL query using a SWbemServices instance and appends the values to dst.
+//
+// dst must have type *[]S or *[]*S, for some struct type S. Fields selected in
+// the query must have the same name in dst. Supported types are all signed and
+// unsigned integers, time.Time, string, bool, or a pointer to one of those.
+// Array types are not supported.
+//
+// By default, the local machine and default namespace are used. These can be
+// changed using connectServerArgs. See
+// http://msdn.microsoft.com/en-us/library/aa393720.aspx for details.
+func (s *SWbemServices) Query(query string, dst interface{}, connectServerArgs ...interface{}) error {
+	s.lQueryorClose.Lock()
+	if s == nil || s.sWbemLocatorIDispatch == nil {
+		s.lQueryorClose.Unlock()
+		return fmt.Errorf("SWbemServices is not Initialized")
+	}
+	if s.queries == nil {
+		s.lQueryorClose.Unlock()
+		return fmt.Errorf("SWbemServices has been closed")
+	}
+
+	//fmt.Println("Query: Sending query request")
+	qr := queryRequest{
+		query:    query,
+		dst:      dst,
+		args:     connectServerArgs,
+		finished: make(chan error),
+	}
+	s.queries <- &qr
+	s.lQueryorClose.Unlock()
+	err, ok := <-qr.finished
+	if ok {
+		//fmt.Println("Query: Finished with error")
+		return err //Send error to caller
+	}
+	//fmt.Println("Query: Finished")
+	return nil
+}
+
+func (s *SWbemServices) queryBackground(q *queryRequest) error {
+	if s == nil || s.sWbemLocatorIDispatch == nil {
+		return fmt.Errorf("SWbemServices is not Initialized")
+	}
+	wmi := s.sWbemLocatorIDispatch //Should just rename in the code, but this will help as we break things apart
+	//fmt.Println("queryBackground: Starting")
+
+	dv := reflect.ValueOf(q.dst)
+	if dv.Kind() != reflect.Ptr || dv.IsNil() {
+		return ErrInvalidEntityType
+	}
+	dv = dv.Elem()
+	mat, elemType := checkMultiArg(dv)
+	if mat == multiArgTypeInvalid {
+		return ErrInvalidEntityType
+	}
+
+	// service is a SWbemServices
+	serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer", q.args...)
+	if err != nil {
+		return err
+	}
+	service := serviceRaw.ToIDispatch()
+	defer serviceRaw.Clear()
+
+	// result is a SWBemObjectSet
+	resultRaw, err := oleutil.CallMethod(service, "ExecQuery", q.query)
+	if err != nil {
+		return err
+	}
+	result := resultRaw.ToIDispatch()
+	defer resultRaw.Clear()
+
+	count, err := oleInt64(result, "Count")
+	if err != nil {
+		return err
+	}
+
+	enumProperty, err := result.GetProperty("_NewEnum")
+	if err != nil {
+		return err
+	}
+	defer enumProperty.Clear()
+
+	enum, err := enumProperty.ToIUnknown().IEnumVARIANT(ole.IID_IEnumVariant)
+	if err != nil {
+		return err
+	}
+	if enum == nil {
+		return fmt.Errorf("can't get IEnumVARIANT, enum is nil")
+	}
+	defer enum.Release()
+
+	// Initialize a slice with Count capacity
+	dv.Set(reflect.MakeSlice(dv.Type(), 0, int(count)))
+
+	var errFieldMismatch error
+	for itemRaw, length, err := enum.Next(1); length > 0; itemRaw, length, err = enum.Next(1) {
+		if err != nil {
+			return err
+		}
+
+		err := func() error {
+			// item is a SWbemObject, but really a Win32_Process
+			item := itemRaw.ToIDispatch()
+			defer item.Release()
+
+			ev := reflect.New(elemType)
+			if err = s.cWMIClient.loadEntity(ev.Interface(), item); err != nil {
+				if _, ok := err.(*ErrFieldMismatch); ok {
+					// We continue loading entities even in the face of field mismatch errors.
+					// If we encounter any other error, that other error is returned. Otherwise,
+					// an ErrFieldMismatch is returned.
+					errFieldMismatch = err
+				} else {
+					return err
+				}
+			}
+			if mat != multiArgTypeStructPtr {
+				ev = ev.Elem()
+			}
+			dv.Set(reflect.Append(dv, ev))
+			return nil
+		}()
+		if err != nil {
+			return err
+		}
+	}
+	//fmt.Println("queryBackground: Finished")
+	return errFieldMismatch
+}

vendor/github.com/StackExchange/wmi/wmi.go

@@ -72,7 +72,10 @@ func QueryNamespace(query string, dst interface{}, namespace string) error {
 //
 // Query is a wrapper around DefaultClient.Query.
 func Query(query string, dst interface{}, connectServerArgs ...interface{}) error {
-	return DefaultClient.Query(query, dst, connectServerArgs...)
+	if DefaultClient.SWbemServicesClient == nil {
+		return DefaultClient.Query(query, dst, connectServerArgs...)
+	}
+	return DefaultClient.SWbemServicesClient.Query(query, dst, connectServerArgs...)
 }
 
 // A Client is an WMI query client.
@@ -99,6 +102,11 @@ type Client struct {
 	// Setting this to true allows custom queries to be used with full
 	// struct definitions instead of having to define multiple structs.
 	AllowMissingFields bool
+
+	// SWbemServiceClient is an optional SWbemServices object that can be
+	// initialized and then reused across multiple queries. If it is null
+	// then the method will initialize a new temporary client each time.
+	SWbemServicesClient *SWbemServices
 }
 
 // DefaultClient is the default Client and is used by Query, QueryNamespace
@@ -362,6 +370,21 @@ func (c *Client) loadEntity(dst interface{}, src *ole.IDispatch) (errFieldMismat
 				}
 			}
 		default:
+			// Only support []string slices for now
+			if f.Kind() == reflect.Slice && f.Type().Elem().Kind() == reflect.String {
+				safeArray := prop.ToArray()
+				if safeArray != nil {
+					arr := safeArray.ToValueArray()
+					fArr := reflect.MakeSlice(f.Type(), len(arr), len(arr))
+					for i, v := range arr {
+						s := fArr.Index(i)
+						s.SetString(v.(string))
+					}
+					f.Set(fArr)
+					break
+				}
+			}
+
 			typeof := reflect.TypeOf(val)
 			if typeof == nil && (isPtr || c.NonePtrZero) {
 				if (isPtr && c.PtrNil) || (!isPtr && c.NonePtrZero) {

vendor/github.com/go-ole/go-ole/variant_s390x.go

@@ -0,0 +1,12 @@
+// +build s390x
+
+package ole
+
+type VARIANT struct {
+	VT         VT      //  2
+	wReserved1 uint16  //  4
+	wReserved2 uint16  //  6
+	wReserved3 uint16  //  8
+	Val        int64   // 16
+	_          [8]byte // 24
+}

vendor/vendor.json

@@ -9,10 +9,10 @@
 			"revisionTime": "2016-08-29T20:23:21Z"
 		},
 		{
-			"checksumSHA1": "9NR0rrcAT5J76C5xMS4AVksS9o0=",
+			"checksumSHA1": "qtjd74+bErubh+qyv3s+lWmn9wc=",
 			"path": "github.com/StackExchange/wmi",
-			"revision": "e54cbda6595d7293a7a468ccf9525f6bc8887f99",
+			"revision": "ea383cf3ba6ec950874b8486cd72356d007c768f",
-			"revisionTime": "2016-08-11T21:45:55Z"
+			"revisionTime": "2017-04-10T19:29:09Z"
 		},
 		{
 			"checksumSHA1": "spyv5/YFBjYyZLZa1U2LBfDR8PM=",
@@ -21,10 +21,10 @@
 			"revisionTime": "2016-08-04T10:47:26Z"
 		},
 		{
-			"checksumSHA1": "wDZdTaY9JiqqqnF4c3pHP71nWmk=",
+			"checksumSHA1": "oyULO1B/l2OLy2xereo+2w3hYk4=",
 			"path": "github.com/go-ole/go-ole",
-			"revision": "7dfdcf409020452e29b4babcbb22f984d2aa308a",
+			"revision": "02d3668a0cf01f58411cc85cd37c174c257ec7c2",
-			"revisionTime": "2016-07-29T03:38:29Z"
+			"revisionTime": "2017-06-01T13:56:11Z"
 		},
 		{
 			"checksumSHA1": "qLYVTQDhgrVIeZ2KI9eZV51mmug=",