diff --git a/internal/iostreams/traffic.go b/internal/iostreams/traffic.go index 93e7d26..65470c2 100644 --- a/internal/iostreams/traffic.go +++ b/internal/iostreams/traffic.go @@ -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 @@ -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 @@ -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 } @@ -87,7 +172,7 @@ 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"]) @@ -95,7 +180,7 @@ func addTrafficObjectFields(obj map[string]any) { } } -func normalizeTrafficSeries(obj map[string]any) { +func normalizeTrafficSeries(obj map[string]any, unit trafficUnit) { columnKey, rowKey := seriesKeys(obj) if columnKey == "" { return @@ -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 { @@ -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 { @@ -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) { @@ -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) } diff --git a/internal/iostreams/traffic_test.go b/internal/iostreams/traffic_test.go index 001a407..e30dd28 100644 --- a/internal/iostreams/traffic_test.go +++ b/internal/iostreams/traffic_test.go @@ -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"]) } } @@ -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:]) } } @@ -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}, @@ -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)) @@ -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()) }