Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 115 additions & 23 deletions internal/iostreams/traffic.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,34 @@ import (
)

const (
// TrafficHumanUnit is the single unit used by structured traffic output.
// TrafficHumanUnit is the default unit kept for compatibility with callers
// that used the original fixed-unit implementation. Structured output now
// selects the unit from trafficUnits based on the largest value.
TrafficHumanUnit = "GiB"
// TrafficUnitSystem describes the byte base used by TrafficHumanUnit.
// TrafficUnitSystem describes the default unit's byte base.
TrafficUnitSystem = "IEC (1 GiB = 1024^3 B)"
// TrafficBytesPerUnit is the number of bytes in one TrafficHumanUnit.
// TrafficBytesPerUnit is the number of bytes in the default unit.
TrafficBytesPerUnit int64 = 1 << 30
)

type trafficUnit struct {
name string
bytesPerUnit int64
power uint
}

var trafficUnits = []trafficUnit{
{name: "B", bytesPerUnit: 1},
{name: "KiB", bytesPerUnit: 1 << 10, power: 1},
{name: "MiB", bytesPerUnit: 1 << 20, power: 2},
{name: "GiB", bytesPerUnit: 1 << 30, power: 3},
{name: "TiB", bytesPerUnit: 1 << 40, power: 4},
{name: "PiB", bytesPerUnit: 1 << 50, power: 5},
{name: "EiB", bytesPerUnit: 1 << 60, power: 6},
}

const trafficHumanPrecision = 3

var trafficRawFields = []struct {
raw string
human string
Expand All @@ -41,29 +61,94 @@ func AddTrafficHumanFields(body []byte) ([]byte, error) {
return nil, err
}

normalizeTrafficValue(payload)
unit := selectTrafficUnit(payload)
normalizeTrafficValue(payload, unit)
return json.Marshal(payload)
}

func normalizeTrafficValue(value any) {
func selectTrafficUnit(value any) trafficUnit {
maxBytes := new(big.Int)
collectTrafficBytes(value, maxBytes)

selected := trafficUnits[0]
for _, candidate := range trafficUnits {
if maxBytes.Cmp(big.NewInt(candidate.bytesPerUnit)) >= 0 {
selected = candidate
}
}
return selected
}

func collectTrafficBytes(value any, maxBytes *big.Int) {
switch v := value.(type) {
case map[string]any:
for _, field := range trafficRawFields {
if raw, ok := v[field.raw]; ok {
updateMaxTrafficBytes(raw, maxBytes)
}
}

// Columnar series keep field names in a separate columns/fields array,
// so inspect those rows explicitly before walking nested values.
if columnKey, rowKey := seriesKeys(v); columnKey != "" {
columns, _ := v[columnKey].([]any)
rows, _ := v[rowKey].([]any)
indexes := make(map[string]int, len(columns))
for i, column := range columns {
if name, ok := column.(string); ok {
indexes[name] = i
}
}
for _, field := range trafficRawFields {
index, ok := indexes[field.raw]
if !ok {
continue
}
for _, row := range rows {
cells, ok := row.([]any)
if ok && index < len(cells) {
updateMaxTrafficBytes(cells[index], maxBytes)
}
}
}
}

for _, child := range v {
collectTrafficBytes(child, maxBytes)
}
case []any:
for _, child := range v {
collectTrafficBytes(child, maxBytes)
}
}
}

func updateMaxTrafficBytes(value any, maxBytes *big.Int) {
bytesValue, ok := trafficBytesValue(value)
if ok && bytesValue.Cmp(maxBytes) > 0 {
maxBytes.Set(bytesValue)
}
}

func normalizeTrafficValue(value any, unit trafficUnit) {
switch v := value.(type) {
case map[string]any:
for _, child := range v {
normalizeTrafficValue(child)
normalizeTrafficValue(child, unit)
}
normalizeTrafficSeries(v)
addTrafficObjectFields(v)
normalizeTrafficSeries(v, unit)
addTrafficObjectFields(v, unit)
if _, ok := v["summary"]; ok {
setTrafficMetadata(v)
setTrafficMetadata(v, unit)
}
case []any:
for _, child := range v {
normalizeTrafficValue(child)
normalizeTrafficValue(child, unit)
}
}
}

func addTrafficObjectFields(obj map[string]any) {
func addTrafficObjectFields(obj map[string]any, unit trafficUnit) {
var (
values = make(map[string]*big.Int, len(trafficRawFields))
matched bool
Expand All @@ -78,7 +163,7 @@ func addTrafficObjectFields(obj map[string]any) {
if !ok {
continue
}
obj[field.human] = formatTrafficGiB(bytesValue)
obj[field.human] = formatTraffic(bytesValue, unit)
values[field.raw] = bytesValue
matched = true
}
Expand All @@ -87,15 +172,15 @@ func addTrafficObjectFields(obj map[string]any) {
return
}

setTrafficMetadata(obj)
setTrafficMetadata(obj, unit)
if len(values) == len(trafficRawFields) {
var parts big.Int
parts.Add(values["rx"], values["tx"])
obj["trafficReconciled"] = parts.Cmp(values["total"]) == 0
}
}

func normalizeTrafficSeries(obj map[string]any) {
func normalizeTrafficSeries(obj map[string]any, unit trafficUnit) {
columnKey, rowKey := seriesKeys(obj)
if columnKey == "" {
return
Expand Down Expand Up @@ -140,7 +225,7 @@ func normalizeTrafficSeries(obj map[string]any) {
}
}

setTrafficMetadata(obj)
setTrafficMetadata(obj, unit)
for rowIndex, row := range rows {
cells, ok := row.([]any)
if !ok {
Expand All @@ -156,7 +241,7 @@ func normalizeTrafficSeries(obj map[string]any) {
continue
}
if bytesValue, ok := trafficBytesValue(cells[rawIndex]); ok {
cells[humanIndex] = formatTrafficGiB(bytesValue)
cells[humanIndex] = formatTraffic(bytesValue, unit)
}
}
if reconciledIndex >= 0 {
Expand All @@ -182,10 +267,10 @@ func normalizeTrafficSeries(obj map[string]any) {
obj[rowKey] = rows
}

func setTrafficMetadata(obj map[string]any) {
obj["trafficUnit"] = TrafficHumanUnit
obj["trafficUnitSystem"] = TrafficUnitSystem
obj["trafficBytesPerUnit"] = TrafficBytesPerUnit
func setTrafficMetadata(obj map[string]any, unit trafficUnit) {
obj["trafficUnit"] = unit.name
obj["trafficUnitSystem"] = unit.system()
obj["trafficBytesPerUnit"] = unit.bytesPerUnit
}

func trafficBytesValue(value any) (*big.Int, bool) {
Expand Down Expand Up @@ -220,8 +305,15 @@ func trafficNumberString(value any) (string, bool) {
}
}

func formatTrafficGiB(bytesValue *big.Int) string {
func (u trafficUnit) system() string {
if u.power == 0 {
return "IEC (bytes)"
}
return fmt.Sprintf("IEC (1 %s = 1024^%d B)", u.name, u.power)
}

func formatTraffic(bytesValue *big.Int, unit trafficUnit) string {
value := new(big.Rat).SetInt(bytesValue)
value.Quo(value, new(big.Rat).SetInt64(TrafficBytesPerUnit))
return fmt.Sprintf("%s %s", value.FloatString(2), TrafficHumanUnit)
value.Quo(value, new(big.Rat).SetInt64(unit.bytesPerUnit))
return fmt.Sprintf("%s %s", value.FloatString(trafficHumanPrecision), unit.name)
}
101 changes: 86 additions & 15 deletions internal/iostreams/traffic_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -41,22 +41,22 @@ func TestAddTrafficHumanFields_Rows(t *testing.T) {
if row["tx"] != float64(312377884) {
t.Errorf("raw tx was changed: %v", row["tx"])
}
if row["txHuman"] != "0.29 GiB" {
t.Errorf("txHuman = %v, want 0.29 GiB", row["txHuman"])
if row["txHuman"] != "0.291 GiB" {
t.Errorf("txHuman = %v, want 0.291 GiB", row["txHuman"])
}
if row["rxHuman"] != "3.04 GiB" {
t.Errorf("rxHuman = %v, want 3.04 GiB", row["rxHuman"])
if row["rxHuman"] != "3.045 GiB" {
t.Errorf("rxHuman = %v, want 3.045 GiB", row["rxHuman"])
}
if row["totalHuman"] != "3.34 GiB" {
t.Errorf("totalHuman = %v, want 3.34 GiB", row["totalHuman"])
if row["totalHuman"] != "3.336 GiB" {
t.Errorf("totalHuman = %v, want 3.336 GiB", row["totalHuman"])
}
if row["trafficReconciled"] != true {
t.Errorf("trafficReconciled = %v, want true", row["trafficReconciled"])
}

trend := payload.Trend[0]
if trend["totalHuman"] != "1.14 GiB" {
t.Errorf("trend totalHuman = %v, want 1.14 GiB", trend["totalHuman"])
if trend["totalHuman"] != "1.137 GiB" {
t.Errorf("trend totalHuman = %v, want 1.137 GiB", trend["totalHuman"])
}
}

Expand Down Expand Up @@ -95,8 +95,8 @@ func TestAddTrafficHumanFields_ColumnarSeries(t *testing.T) {
t.Fatalf("unexpected data shape: %#v", series.Data)
}
row := series.Data[0]
if row[4] != "0.29 GiB" || row[5] != "3.04 GiB" || row[6] != "3.34 GiB" || row[7] != true {
t.Errorf("human cells = %v, want [0.29 GiB 3.04 GiB 3.34 GiB]", row[4:])
if row[4] != "0.291 GiB" || row[5] != "3.045 GiB" || row[6] != "3.336 GiB" || row[7] != true {
t.Errorf("human cells = %v, want [0.291 GiB 3.045 GiB 3.336 GiB]", row[4:])
}
}

Expand All @@ -106,6 +106,77 @@ func TestAddTrafficHumanFields_InvalidJSON(t *testing.T) {
}
}

func TestSelectTrafficUnit(t *testing.T) {
tests := []struct {
name string
bytes int64
want string
}{
{name: "bytes", bytes: 1023, want: "B"},
{name: "kib", bytes: 1 << 10, want: "KiB"},
{name: "mib", bytes: 1 << 20, want: "MiB"},
{name: "gib", bytes: 1 << 30, want: "GiB"},
{name: "tib", bytes: 1 << 40, want: "TiB"},
{name: "pib", bytes: 1 << 50, want: "PiB"},
{name: "eib", bytes: 1 << 60, want: "EiB"},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := selectTrafficUnit(map[string]any{"total": tt.bytes})
if got.name != tt.want {
t.Errorf("unit = %q, want %q", got.name, tt.want)
}
})
}
}

func TestAddTrafficHumanFields_AdaptiveUnitIsShared(t *testing.T) {
input := []byte(`{
"summary":[
{"tx":4434710,"rx":4466763,"total":8901473},
{"tx":25670393,"rx":141266066,"total":166936459}
]
}`)

got, err := AddTrafficHumanFields(input)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}

var payload struct {
TrafficUnit string `json:"trafficUnit"`
TrafficUnitSystem string `json:"trafficUnitSystem"`
TrafficBytesPerUnit int64 `json:"trafficBytesPerUnit"`
Summary []map[string]interface{} `json:"summary"`
}
if err := json.Unmarshal(got, &payload); err != nil {
t.Fatalf("invalid output: %v", err)
}

if payload.TrafficUnit != "MiB" {
t.Errorf("trafficUnit = %q, want MiB", payload.TrafficUnit)
}
if payload.TrafficUnitSystem != "IEC (1 MiB = 1024^2 B)" {
t.Errorf("trafficUnitSystem = %q", payload.TrafficUnitSystem)
}
if payload.TrafficBytesPerUnit != 1<<20 {
t.Errorf("trafficBytesPerUnit = %d, want %d", payload.TrafficBytesPerUnit, 1<<20)
}
if got := payload.Summary[0]["txHuman"]; got != "4.229 MiB" {
t.Errorf("txHuman = %v, want 4.229 MiB", got)
}
if got := payload.Summary[0]["rxHuman"]; got != "4.260 MiB" {
t.Errorf("rxHuman = %v, want 4.260 MiB", got)
}
if got := payload.Summary[0]["totalHuman"]; got != "8.489 MiB" {
t.Errorf("totalHuman = %v, want 8.489 MiB", got)
}
if got := payload.Summary[1]["totalHuman"]; got != "159.203 MiB" {
t.Errorf("second totalHuman = %v, want 159.203 MiB", got)
}
}

func TestAddTrafficHumanFields_IM3173Samples(t *testing.T) {
input := []byte(`{"rows":[
{"date":"08-06","tx":312377884,"rx":3269296004,"total":3581673888},
Expand All @@ -129,10 +200,10 @@ func TestAddTrafficHumanFields_IM3173Samples(t *testing.T) {
date string
tx, rx, total string
}{
{date: "08-06", tx: "0.29 GiB", rx: "3.04 GiB", total: "3.34 GiB"},
{date: "08-07", tx: "0.27 GiB", rx: "0.87 GiB", total: "1.14 GiB"},
{date: "08-13", tx: "0.06 GiB", rx: "0.67 GiB", total: "0.73 GiB"},
{date: "summary", tx: "1.21 GiB", rx: "10.52 GiB", total: "11.73 GiB"},
{date: "08-06", tx: "0.291 GiB", rx: "3.045 GiB", total: "3.336 GiB"},
{date: "08-07", tx: "0.266 GiB", rx: "0.872 GiB", total: "1.137 GiB"},
{date: "08-13", tx: "0.056 GiB", rx: "0.671 GiB", total: "0.728 GiB"},
{date: "summary", tx: "1.210 GiB", rx: "10.515 GiB", total: "11.725 GiB"},
}
if len(payload.Rows) != len(want) {
t.Fatalf("got %d rows, want %d", len(payload.Rows), len(want))
Expand All @@ -156,7 +227,7 @@ func TestTrafficHumanFieldsStructuredOnly(t *testing.T) {
if err := FormatOutput(body, jsonIO, "json", WithStructuredTransform(AddTrafficHumanFields)); err != nil {
t.Fatalf("json output: %v", err)
}
if !strings.Contains(jsonOut.String(), `"totalHuman"`) || !strings.Contains(jsonOut.String(), `3.34 GiB`) {
if !strings.Contains(jsonOut.String(), `"totalHuman"`) || !strings.Contains(jsonOut.String(), `3.336 GiB`) {
t.Errorf("structured output missing totalHuman: %s", jsonOut.String())
}

Expand Down