From b93c4acbc4c8fdbd4f8fd0b3ae20d53425eeceb9 Mon Sep 17 00:00:00 2001
From: Zain Budhwani <99770260+zbud-msft@users.noreply.github.com>
Date: Wed, 1 Mar 2023 15:45:43 -0800
Subject: [PATCH] Fix crash when retrieving cpu utilization (#70) (#71)

Fix crash in which there was a divide by zero error inside retrieving cpu utilization
---
 gnmi_server/server_test.go         | 81 ++++++++++++++++++++++++++++++
 go.mod                             |  1 +
 sonic_data_client/non_db_client.go | 26 ++++++----
 3 files changed, 98 insertions(+), 10 deletions(-)

diff --git a/gnmi_server/server_test.go b/gnmi_server/server_test.go
index 79205fe2..7eb51771 100644
--- a/gnmi_server/server_test.go
+++ b/gnmi_server/server_test.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"flag"
 	"fmt"
+        "sync"
 	"strings"
 
 	testcert "github.com/Azure/sonic-telemetry/testdata/tls"
@@ -39,9 +40,11 @@ import (
 	sgpb "github.com/Azure/sonic-telemetry/proto/gnoi"
 	sdc "github.com/Azure/sonic-telemetry/sonic_data_client"
 	sdcfg "github.com/Azure/sonic-telemetry/sonic_db_config"
+        linuxproc "github.com/c9s/goprocinfo/linux"
 	gclient "github.com/jipanyang/gnmi/client/gnmi"
 	"github.com/jipanyang/gnxi/utils/xpath"
 	gnoi_system_pb "github.com/openconfig/gnoi/system"
+	"github.com/agiledragon/gomonkey/v2"
 )
 
 var clientTypes = []string{gclient.Type}
@@ -2303,6 +2306,84 @@ func TestGNOI(t *testing.T) {
     }
 }
 
+func TestCPUUtilization(t *testing.T) {
+    mock := gomonkey.ApplyFunc(sdc.PollStats, func() {
+	var i uint64
+	for i = 0; i < 3000; i++ {
+		sdc.WriteStatsToBuffer(&linuxproc.Stat{})
+	}
+    })
+
+    defer mock.Reset()
+    s := createServer(t, 8081)
+    go runServer(t, s)
+    defer s.s.Stop()
+
+    tests := []struct {
+        desc    string
+	q       client.Query
+	want    []client.Notification
+	poll    int
+    }{
+        {
+            desc: "poll query for CPU Utilization",
+	    poll: 10,
+	    q: client.Query{
+                Target: "OTHERS",
+		Type:    client.Poll,
+		Queries: []client.Path{{"platform", "cpu"}},
+		TLS:     &tls.Config{InsecureSkipVerify: true},
+	    },
+	    want: []client.Notification{
+                client.Connected{},
+		client.Sync{},
+	    },
+        },
+    }
+
+    for _, tt := range tests {
+        t.Run(tt.desc, func(t *testing.T) {
+            q := tt.q
+	    q.Addrs = []string{"127.0.0.1:8081"}
+            c := client.New()
+            var gotNoti []client.Notification
+            q.NotificationHandler = func(n client.Notification) error {
+                if nn, ok := n.(client.Update); ok {
+                    nn.TS = time.Unix(0, 200)
+		    gotNoti = append(gotNoti, nn)
+                } else {
+                    gotNoti = append(gotNoti, n)
+	        }
+                return nil
+            }
+
+            wg := new(sync.WaitGroup)
+            wg.Add(1)
+
+            go func() {
+                defer wg.Done()
+                if err := c.Subscribe(context.Background(), q); err != nil {
+                    t.Errorf("c.Subscribe(): got error %v, expected nil", err)
+                }
+            }()
+
+            wg.Wait()
+
+            for i := 0; i < tt.poll; i++ {
+                if err := c.Poll(); err != nil {
+                    t.Errorf("c.Poll(): got error %v, expected nil", err)
+                }
+	    }
+
+            if len(gotNoti) == 0 {
+                t.Errorf("expected non zero notifications")
+            }
+
+            c.Close()
+        })
+    }
+}
+
 func TestBundleVersion(t *testing.T) {
     s := createServer(t, 8087)
     go runServer(t, s)
diff --git a/go.mod b/go.mod
index 87f3c6fc..e7c48a68 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.12
 require (
 	github.com/Azure/sonic-mgmt-common v0.0.0-00010101000000-000000000000
 	github.com/Workiva/go-datastructures v1.0.50
+	github.com/agiledragon/gomonkey/v2 v2.8.0
 	github.com/c9s/goprocinfo v0.0.0-20191125144613-4acdd056c72d
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/go-redis/redis v6.15.6+incompatible
diff --git a/sonic_data_client/non_db_client.go b/sonic_data_client/non_db_client.go
index 12336a68..422196f4 100644
--- a/sonic_data_client/non_db_client.go
+++ b/sonic_data_client/non_db_client.go
@@ -132,10 +132,13 @@ type cpuUtil struct {
 }
 
 func getCpuUtilPercents(cur, last *linuxproc.CPUStat) uint64 {
-	curTotal := (cur.User + cur.Nice + cur.System + cur.Idle + cur.IOWait + cur.IRQ + cur.SoftIRQ + cur.Steal + cur.Guest + cur.GuestNice)
-	lastTotal := (last.User + last.Nice + last.System + last.Idle + last.IOWait + last.IRQ + last.SoftIRQ + last.Steal + last.Guest + last.GuestNice)
+	curTotal := (cur.User + cur.Nice + cur.System + cur.Idle + cur.IOWait + cur.IRQ + cur.SoftIRQ + cur.Steal)
+	lastTotal := (last.User + last.Nice + last.System + last.Idle + last.IOWait + last.IRQ + last.SoftIRQ + last.Steal)
 	idleTicks := cur.Idle - last.Idle
 	totalTicks := curTotal - lastTotal
+	if totalTicks == 0 { // No change in CPU Utilization
+            return 0
+	}
 	return 100 * (totalTicks - idleTicks) / totalTicks
 }
 
@@ -332,7 +335,15 @@ func getBuildVersion() ([]byte, error) {
 	return b, nil
 }
 
-func pollStats() {
+func WriteStatsToBuffer(stat *linuxproc.Stat) {
+	statsR.mu.Lock()
+	statsR.buff[statsR.writeIdx] = stat
+	statsR.writeIdx++
+	statsR.writeIdx %= statsRingCap
+	statsR.mu.Unlock()
+}
+
+func PollStats() {
 	for {
 		stat, err := linuxproc.ReadStat("/proc/stat")
 		if err != nil {
@@ -340,12 +351,7 @@ func pollStats() {
 			continue
 		}
 
-		statsR.mu.Lock()
-
-		statsR.buff[statsR.writeIdx] = stat
-		statsR.writeIdx++
-		statsR.writeIdx %= statsRingCap
-		statsR.mu.Unlock()
+		WriteStatsToBuffer(stat)
 		time.Sleep(time.Millisecond * 100)
 	}
 
@@ -355,7 +361,7 @@ func init() {
 	clientTrie = NewTrie()
 	clientTrie.clientTriePopulate()
 	statsR.buff = make([]*linuxproc.Stat, statsRingCap)
-	go pollStats()
+	go PollStats()
 }
 
 type NonDbClient struct {