From 80461345049ce26f51370e8652400fa2e3c3709e Mon Sep 17 00:00:00 2001
From: Ryan Fitzpatrick <rmfitzpatrick@signalfx.com>
Date: Tue, 6 Oct 2020 17:03:49 +0000
Subject: [PATCH] Fix VirtualizationWithContext() race in linux

---
 host/host_test.go               | 26 +++++++++++++++------
 internal/common/common_linux.go | 52 ++++++++++++++++++++++++++++++++---------
 2 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/host/host_test.go b/host/host_test.go
index 0970d21..e1749f3 100644
--- a/host/host_test.go
+++ b/host/host_test.go
@@ -3,6 +3,7 @@ package host
 import (
 	"fmt"
 	"os"
+	"sync"
 	"testing"
 
 	"github.com/shirou/gopsutil/internal/common"
@@ -144,13 +145,24 @@ func TestTemperatureStat_String(t *testing.T) {
 }
 
 func TestVirtualization(t *testing.T) {
-	system, role, err := Virtualization()
-	skipIfNotImplementedErr(t, err)
-	if err != nil {
-		t.Errorf("Virtualization() failed, %v", err)
-	}
-
-	t.Logf("Virtualization(): %s, %s", system, role)
+	wg := sync.WaitGroup{}
+	testCount := 10
+	wg.Add(testCount)
+	for i := 0; i < testCount; i++ {
+		go func(j int) {
+			system, role, err := Virtualization()
+			wg.Done()
+			skipIfNotImplementedErr(t, err)
+			if err != nil {
+				t.Errorf("Virtualization() failed, %v", err)
+			}
+
+			if j == 9 {
+				t.Logf("Virtualization(): %s, %s", system, role)
+			}
+		}(i)
+	}
+	wg.Wait()
 }
 
 func TestKernelVersion(t *testing.T) {
diff --git a/internal/common/common_linux.go b/internal/common/common_linux.go
index 62fac5e..6889e91 100644
--- a/internal/common/common_linux.go
+++ b/internal/common/common_linux.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -110,16 +111,49 @@ func Virtualization() (string, string, error) {
 	return VirtualizationWithContext(context.Background())
 }
 
-var virtualizationCache map[string]string
+type virtCache struct {
+	cache map[string]string
+	lock sync.Mutex
+	ok bool
+}
+
+func (v *virtCache) setValue(system, role string) {
+	v.lock.Lock()
+	defer v.lock.Unlock()
+	v.cache["system"] = system
+	v.cache["role"] = role
+	v.ok = true
+}
+
+func (v *virtCache) getValue() (string, string, bool) {
+	v.lock.Lock()
+	defer v.lock.Unlock()
+	return v.cache["system"], v.cache["role"], v.ok
+}
+
+var (
+	once sync.Once
+	virtualization *virtCache
+)
+
+func virtualizationCache() *virtCache {
+	once.Do(func() {
+		virtualization = &virtCache{
+			cache: make(map[string]string),
+			lock: sync.Mutex{},
+			ok: false,
+		}
+	})
+
+	return virtualization
+}
 
 func VirtualizationWithContext(ctx context.Context) (string, string, error) {
 	// if cached already, return from cache
-	if virtualizationCache != nil {
-		return virtualizationCache["system"], virtualizationCache["role"], nil
+	system, role, ok := virtualizationCache().getValue()
+	if ok {
+		return system, role, nil
 	}
-	
-	var system string
-	var role string
 
 	filename := HostProc("xen")
 	if PathExists(filename) {
@@ -240,11 +274,7 @@ func VirtualizationWithContext(ctx context.Context) (string, string, error) {
 	}
 	
 	// before returning for the first time, cache the system and role
-	virtualizationCache = map[string]string{
-		"system":	system,
-		"role": 	role,
-	}
-	
+	virtualizationCache().setValue(system, role)
 	return system, role, nil
 }