From 175c54acf13aee43592120ad0567257f01684dc1 Mon Sep 17 00:00:00 2001 From: Calle Pettersson Date: Mon, 26 Jun 2017 21:03:17 +0200 Subject: [PATCH] Fix memory leak by updating go-ole and adding wbem initialization (#82) --- exporter.go | 15 + vendor/github.com/StackExchange/wmi/README.md | 4 +- .../StackExchange/wmi/swbemservices.go | 260 ++++++++++++++++++ vendor/github.com/StackExchange/wmi/wmi.go | 25 +- .../github.com/go-ole/go-ole/variant_s390x.go | 12 + vendor/vendor.json | 12 +- 6 files changed, 320 insertions(+), 8 deletions(-) create mode 100644 vendor/github.com/StackExchange/wmi/swbemservices.go create mode 100644 vendor/github.com/go-ole/go-ole/variant_s390x.go diff --git a/exporter.go b/exporter.go index 0d6b88fd..bc2008d2 100644 --- a/exporter.go +++ b/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) diff --git a/vendor/github.com/StackExchange/wmi/README.md b/vendor/github.com/StackExchange/wmi/README.md index 3d5f67e1..426d1a46 100644 --- a/vendor/github.com/StackExchange/wmi/README.md +++ b/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. diff --git a/vendor/github.com/StackExchange/wmi/swbemservices.go b/vendor/github.com/StackExchange/wmi/swbemservices.go new file mode 100644 index 00000000..9765a53f --- /dev/null +++ b/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 +} diff --git a/vendor/github.com/StackExchange/wmi/wmi.go b/vendor/github.com/StackExchange/wmi/wmi.go index 5cb56ff1..9c688b03 100644 --- a/vendor/github.com/StackExchange/wmi/wmi.go +++ b/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) { diff --git a/vendor/github.com/go-ole/go-ole/variant_s390x.go b/vendor/github.com/go-ole/go-ole/variant_s390x.go new file mode 100644 index 00000000..9874ca66 --- /dev/null +++ b/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 +} diff --git a/vendor/vendor.json b/vendor/vendor.json index 6ece5658..22993f6f 100644 --- a/vendor/vendor.json +++ b/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", - "revisionTime": "2016-08-11T21:45:55Z" + "revision": "ea383cf3ba6ec950874b8486cd72356d007c768f", + "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", - "revisionTime": "2016-07-29T03:38:29Z" + "revision": "02d3668a0cf01f58411cc85cd37c174c257ec7c2", + "revisionTime": "2017-06-01T13:56:11Z" }, { "checksumSHA1": "qLYVTQDhgrVIeZ2KI9eZV51mmug=",