From c412fc746e7fcf7d741be892c1432fcb83cabeda Mon Sep 17 00:00:00 2001 From: Mihail Stoykov Date: Wed, 2 Sep 2020 12:53:05 +0300 Subject: [PATCH] Add (Un)Marshal functions to stats.TagSet (#1604) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Also, sort the keys before marshaling them with SystemTagSet for consistent output. fixes #1602 Co-authored-by: Ivan Mirić --- stats/system_tag.go | 56 ++++++++++++++++++++++++++++++++--- stats/system_tag_test.go | 63 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 111 insertions(+), 8 deletions(-) diff --git a/stats/system_tag.go b/stats/system_tag.go index 9d98f64a366..48000bdfecb 100644 --- a/stats/system_tag.go +++ b/stats/system_tag.go @@ -23,6 +23,7 @@ package stats import ( "bytes" "encoding/json" + "sort" "strings" ) @@ -32,12 +33,57 @@ import ( type SystemTagSet uint32 // TagSet is a string to bool map (for lookup efficiency) that is used to keep track -// which system tags and non-system tags should be included with with metrics. +// of which system tags and non-system tags to include. type TagSet map[string]bool -//nolint: golint +// UnmarshalText converts the tag list to TagSet. +func (i *TagSet) UnmarshalText(data []byte) error { + list := bytes.Split(data, []byte(",")) + if *i == nil { + *i = make(TagSet, len(list)) + } + + for _, key := range list { + key := strings.TrimSpace(string(key)) + if key == "" { + continue + } + (*i)[key] = true + } + + return nil +} + +// MarshalJSON converts the TagSet to a list (JS array). +func (i *TagSet) MarshalJSON() ([]byte, error) { + var tags []string + if *i != nil { + tags = make([]string, 0, len(*i)) + for tag := range *i { + tags = append(tags, tag) + } + sort.Strings(tags) + } + + return json.Marshal(tags) +} + +// UnmarshalJSON converts the tag list back to expected tag set. +func (i *TagSet) UnmarshalJSON(data []byte) error { + var tags []string + if err := json.Unmarshal(data, &tags); err != nil { + return err + } + *i = make(TagSet, len(tags)) + for _, tag := range tags { + (*i)[tag] = true + } + + return nil +} + +// Default system tags includes all of the system tags emitted with metrics by default. const ( - // Default system tags includes all of the system tags emitted with metrics by default. TagProto SystemTagSet = 1 << iota TagSubproto TagStatus @@ -131,6 +177,8 @@ func (i *SystemTagSet) MarshalJSON() ([]byte, error) { tags = append(tags, tag.String()) } } + sort.Strings(tags) + return json.Marshal(tags) } @@ -149,7 +197,7 @@ func (i *SystemTagSet) UnmarshalJSON(data []byte) error { // UnmarshalText converts the tag list to SystemTagSet. func (i *SystemTagSet) UnmarshalText(data []byte) error { - var list = bytes.Split(data, []byte(",")) + list := bytes.Split(data, []byte(",")) for _, key := range list { key := strings.TrimSpace(string(key)) diff --git a/stats/system_tag_test.go b/stats/system_tag_test.go index 27c98e3fd17..86f91e53974 100644 --- a/stats/system_tag_test.go +++ b/stats/system_tag_test.go @@ -29,11 +29,12 @@ import ( ) func TestSystemTagSetMarshalJSON(t *testing.T) { - var tests = []struct { + tests := []struct { tagset SystemTagSet expected string }{ {TagIP, `["ip"]`}, + {TagIP | TagProto | TagGroup, `["group","ip","proto"]`}, {0, `null`}, } @@ -46,7 +47,7 @@ func TestSystemTagSetMarshalJSON(t *testing.T) { } func TestSystemTagSet_UnmarshalJSON(t *testing.T) { - var tests = []struct { + tests := []struct { tags []byte sets []SystemTagSet }{ @@ -64,7 +65,7 @@ func TestSystemTagSet_UnmarshalJSON(t *testing.T) { } func TestSystemTagSetTextUnmarshal(t *testing.T) { - var testMatrix = map[string]SystemTagSet{ + testMatrix := map[string]SystemTagSet{ "": 0, "ip": TagIP, "ip,proto": TagIP | TagProto, @@ -74,7 +75,61 @@ func TestSystemTagSetTextUnmarshal(t *testing.T) { } for input, expected := range testMatrix { - var set = new(SystemTagSet) + set := new(SystemTagSet) + err := set.UnmarshalText([]byte(input)) + require.NoError(t, err) + require.Equal(t, expected, *set) + } +} + +func TestTagSetMarshalJSON(t *testing.T) { + tests := []struct { + tagset TagSet + expected string + }{ + {tagset: TagSet{"ip": true, "proto": true, "group": true, "custom": true}, expected: `["custom","group","ip","proto"]`}, + {tagset: TagSet{}, expected: `[]`}, + } + + for _, tc := range tests { + ts := &tc.tagset + got, err := json.Marshal(ts) + require.Nil(t, err) + require.Equal(t, tc.expected, string(got)) + } +} + +func TestTagSet_UnmarshalJSON(t *testing.T) { + tests := []struct { + tags []byte + sets TagSet + }{ + {[]byte(`[]`), TagSet{}}, + {[]byte(`["ip","custom", "proto"]`), TagSet{"ip": true, "proto": true, "custom": true}}, + } + + for _, tc := range tests { + ts := new(TagSet) + require.Nil(t, json.Unmarshal(tc.tags, ts)) + for tag := range tc.sets { + assert.True(t, (*ts)[tag]) + } + } +} + +func TestTagSetTextUnmarshal(t *testing.T) { + testMatrix := map[string]TagSet{ + "": make(TagSet), + "ip": {"ip": true}, + "ip,proto": {"ip": true, "proto": true}, + " ip , proto ": {"ip": true, "proto": true}, + " ip , , proto ": {"ip": true, "proto": true}, + " ip ,, proto ,,": {"ip": true, "proto": true}, + " ip ,custom, proto ,,": {"ip": true, "custom": true, "proto": true}, + } + + for input, expected := range testMatrix { + set := new(TagSet) err := set.UnmarshalText([]byte(input)) require.NoError(t, err) require.Equal(t, expected, *set)