Skip to content
This repository has been archived by the owner on Jan 22, 2022. It is now read-only.

Commit

Permalink
implement requested changes. Also add benchmark tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@gbrayut
gbrayut committed Feb 14, 2017
1 parent f03346a commit 2f76343
Show file tree
Hide file tree
Showing 2 changed files with 105 additions and 88 deletions.
147 changes: 60 additions & 87 deletions swbemservices.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@ package wmi

import (
"fmt"
"github.com/go-ole/go-ole"
"github.com/go-ole/go-ole/oleutil"
"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
Expand All @@ -18,81 +20,64 @@ type SWbemServices struct {
sWbemLocatorIDispatch *ole.IDispatch
queries chan *queryRequest
closeError chan error
lQueryorClose sync.Mutex
}

type queryRequest struct {
query string
dst interface{}
args 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")
//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)

WaitForProcessSignal:
for {
select {
case err, ok := <-initError:
if ok {
return nil, err //Send error to caller
}
break WaitForProcessSignal
default:
//Wait for process method to return an error or signal it is ready (by closing channel)
}
err, ok := <-initError
if ok {
return nil, err //Send error to caller
}
fmt.Println("InitializeSWbemServices: Finished")
//fmt.Println("InitializeSWbemServices: Finished")
return s, nil
}

// Close will clear and release all of the SWbemServices resources
func (s *SWbemServices) Close() error {
fmt.Println("Close: sending close request")
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
WaitForClosedSignal:
for {
select {
case err, ok := <-ce:
if ok {
result = err
}
break WaitForClosedSignal
default:
//Wait for close method to return an error or signal it is finished by closing channel
}
s.lQueryorClose.Unlock()
err, ok := <-ce
if ok {
result = err
}
fmt.Println("Close: finished")
//fmt.Println("Close: finished")
return result
}

// This method will be called on the background thread and must clean up all WMI/OLE and chan resources
func (s *SWbemServices) closeBackground() error {
//TODO: detect errors and return them to the caller somehow? The caller could be InitializeSWbemServices, Query, or Close
fmt.Println("closeBackground: Starting")
s.queries = nil //set channel to nil so we know it is closed
s.sWbemLocatorIDispatch.Release()
s.sWbemLocatorIUnknown.Release()
ole.CoUninitialize()
runtime.UnlockOSThread()
fmt.Println("closeBackground: Finished")
return nil
}

func (s *SWbemServices) process(initError chan error) {
fmt.Println("process: starting background thread initialization")
//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 s.closeBackground() //Does this run on the same locked thread? Can't rely on s.closing as select hasn't started yet
defer runtime.LockOSThread()

err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED)
if err != nil {
Expand All @@ -102,6 +87,7 @@ func (s *SWbemServices) process(initError chan error) {
return
}
}
defer ole.CoUninitialize()

unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
if err != nil {
Expand All @@ -111,46 +97,38 @@ func (s *SWbemServices) process(initError chan error) {
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")
//fmt.Println("process: initialized. closing initError")
close(initError)
fmt.Println("process: waiting for queries")
for {
select {
case q, ok := <-s.queries:
if ok {
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)
} else {
fmt.Println("process: queries channel closed")
errClose := s.closeBackground()
if errClose == nil && s.closeError != nil {
s.closeError <- errClose
}
if s.closeError != nil {
close(s.closeError)
}
}
default:
//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 and appends the values to dst.
// 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
Expand All @@ -161,36 +139,31 @@ func (s *SWbemServices) process(initError chan error) {
// 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")
//fmt.Println("Query: Sending query request")
qr := queryRequest{
query: query,
dst: dst,
args: connectServerArgs,
finished: make(chan error),
}
s.queries <- &qr

WaitForFinishedSignal:
for {
select {
case err, ok := <-qr.finished:
if ok {
fmt.Println("Query: Finished with error")
return err //Send error to caller
}
break WaitForFinishedSignal
default:
//Wait for query to return an error or signal it is finished by closing the channel
}
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")
//fmt.Println("Query: Finished")
return nil
}

Expand All @@ -199,7 +172,7 @@ func (s *SWbemServices) queryBackground(q *queryRequest) error {
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")
//fmt.Println("queryBackground: Starting")

dv := reflect.ValueOf(q.dst)
if dv.Kind() != reflect.Ptr || dv.IsNil() {
Expand All @@ -212,7 +185,7 @@ func (s *SWbemServices) queryBackground(q *queryRequest) error {
}

// service is a SWbemServices
serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer", nil) // todo: use q.args...)
serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer", q.args...)
if err != nil {
return err
}
Expand Down Expand Up @@ -282,6 +255,6 @@ func (s *SWbemServices) queryBackground(q *queryRequest) error {
return err
}
}
fmt.Println("queryBackground: Finished")
//fmt.Println("queryBackground: Finished")
return errFieldMismatch
}
46 changes: 45 additions & 1 deletion swbemservices_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ func TestWbemMemory(t *testing.T) {
var privateMB, allocMB, allocTotalMB float64
for i := 0; i < limit; i++ {
privateMB, allocMB, allocTotalMB = WbemGetMemoryUsageMB(s)
if i%1000 == 0 {
if i%100 == 0 {
privateMB, allocMB, allocTotalMB = WbemGetMemoryUsageMB(s)
fmt.Printf("Time: %4ds Count: %5d Private Memory: %5.1fMB MemStats.Alloc: %4.1fMB MemStats.TotalAlloc: %5.1fMB\n", time.Now().Sub(start)/time.Second, i, privateMB, allocMB, allocTotalMB)
}
Expand All @@ -78,3 +78,47 @@ func WbemGetMemoryUsageMB(s *SWbemServices) (float64, float64, float64) {
}
return float64(dstGetMemoryUsageMB[0].WorkingSetPrivate) / MB, float64(mMemoryUsageMB.Alloc) / MB, float64(mMemoryUsageMB.TotalAlloc) / MB
}

//Run all benchmarks (should run for at least 60s to get a stable number):
//go test -run=NONE -bench=Version -benchtime=120s

//Individual benchmarks:
//go test -run=NONE -bench=NewVersion -benchtime=120s
func BenchmarkNewVersion(b *testing.B) {
s, err := InitializeSWbemServices(DefaultClient)
if err != nil {
b.Fatalf("InitializeSWbemServices: %s", err)
}
var dst []Win32_OperatingSystem
q := CreateQuery(&dst, "")
for n := 0; n < b.N; n++ {
errQuery := s.Query(q, &dst)
if errQuery != nil {
b.Fatalf("Query%d: %s", n, errQuery)
}
count := len(dst)
if count < 1 {
b.Fatalf("Query%d: no results found for Win32_OperatingSystem", n)
}
}
errClose := s.Close()
if errClose != nil {
b.Fatalf("Close: %s", errClose)
}
}

//go test -run=NONE -bench=OldVersion -benchtime=120s
func BenchmarkOldVersion(b *testing.B) {
var dst []Win32_OperatingSystem
q := CreateQuery(&dst, "")
for n := 0; n < b.N; n++ {
errQuery := Query(q, &dst)
if errQuery != nil {
b.Fatalf("Query%d: %s", n, errQuery)
}
count := len(dst)
if count < 1 {
b.Fatalf("Query%d: no results found for Win32_OperatingSystem", n)
}
}
}

0 comments on commit 2f76343

Please sign in to comment.