pkg/tcpip: impl json.Marshaler/Unmarshaler for StatCounter to Marshal/Unmarshal *Stats

Author: ruokeqx

pkg/tcpip/tcpip.go

@@ -30,6 +30,7 @@ package tcpip
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -1633,6 +1634,21 @@ func (s *StatCounter) String() string {
 	return strconv.FormatUint(s.Value(), 10)
 }
 
+// MarshalJSON implements json.Marshaler.MarshalJSON
+func (s *StatCounter) MarshalJSON() ([]byte, error) {
+	return json.Marshal(s.Value())
+}
+
+// UnmarshalJSON implements json.Unmarshaler.UnmarshalJSON
+func (s *StatCounter) UnmarshalJSON(data []byte) error {
+	var val uint64
+	if err := json.Unmarshal(data, &val); err != nil {
+		return err
+	}
+	s.count.Store(val)
+	return nil
+}
+
 // A MultiCounterStat keeps track of two counters at once.
 //
 // +stateify savable

pkg/tcpip/tcpip_test.go

@@ -16,6 +16,7 @@ package tcpip
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net"
@@ -345,3 +346,25 @@ func padTo4(partial string) []byte {
 	}
 	return []byte(partial)
 }
+
+func TestStats_Marshal(t *testing.T) {
+	s := Stats{}.FillIn()
+	s.TCP.ForwardMaxInFlightDrop.Increment()
+	jb, err := json.Marshal(s)
+	if err != nil {
+		t.Fail()
+	}
+	if !bytes.Contains(jb, []byte(`"ForwardMaxInFlightDrop":1`)) {
+		t.Fatalf("Marshal did not contain ForwardMaxInFlightDrop")
+	}
+
+	todo := `{"TCP":{"ForwardMaxInFlightDrop":1}}`
+	var to Stats
+	err = json.Unmarshal([]byte(todo), &to)
+	if err != nil {
+		t.Fail()
+	}
+	if got := to.TCP.ForwardMaxInFlightDrop.Value(); got != uint64(1) {
+		t.Fatalf("got ForwardMaxInFlightDrop.Value() = %d, want = %d", got, uint64(1))
+	}
+}