From 2b972c4f5582b362294449617264fe7899f9290a Mon Sep 17 00:00:00 2001 From: roroghost17 Date: Tue, 18 Aug 2026 19:14:37 +0530 Subject: [PATCH] perf: update json encode usage with sonic marshal --- plugins/logging/main.go | 7 + transports/bifrost-http/handlers/utils.go | 15 ++- .../bifrost-http/handlers/utils_test.go | 120 ++++++++++-------- 3 files changed, 83 insertions(+), 59 deletions(-) diff --git a/plugins/logging/main.go b/plugins/logging/main.go index 418dbb05653..e318874ff2e 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -1751,6 +1751,13 @@ func (p *LoggerPlugin) Inject(_ context.Context, trace *schemas.Trace) error { for i, entry := range pending.entries { entry.PluginLogs = pluginLogsJSON if i == stampIdx { + // Clear both provisional components before backfilling. A partial root + // span (only one attribute present) would otherwise leave the other as + // a stale PostLLMHook estimate, mixing the breakdown's two sources. + if upOK || ovOK { + entry.UpstreamLatency = nil + entry.OverheadLatency = nil + } if upOK { u := upstreamMs entry.UpstreamLatency = &u diff --git a/transports/bifrost-http/handlers/utils.go b/transports/bifrost-http/handlers/utils.go index 20d236de6e8..e67a97ff85f 100644 --- a/transports/bifrost-http/handlers/utils.go +++ b/transports/bifrost-http/handlers/utils.go @@ -12,6 +12,7 @@ import ( "sync" "time" + "github.com/bytedance/sonic" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/transports/bifrost-http/lib" @@ -82,23 +83,29 @@ func IsUniqueConstraintError(err error, identifiers ...string) bool { return false } -// SendJSON sends a JSON response with 200 OK status +// SendJSON sends a JSON response with 200 OK status. func SendJSON(ctx *fasthttp.RequestCtx, data interface{}) { ctx.SetContentType("application/json") - if err := json.NewEncoder(ctx).Encode(data); err != nil { + body, err := sonic.ConfigStd.Marshal(data) + if err != nil { logger.Warn(fmt.Sprintf("Failed to encode JSON response: %v", err)) SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to encode response: %v", err)) + return } + ctx.SetBody(append(body, '\n')) } -// SendJSONWithStatus sends a JSON response with a custom status code +// SendJSONWithStatus sends a JSON response with a custom status code. func SendJSONWithStatus(ctx *fasthttp.RequestCtx, data interface{}, statusCode int) { ctx.SetContentType("application/json") ctx.SetStatusCode(statusCode) - if err := json.NewEncoder(ctx).Encode(data); err != nil { + body, err := sonic.ConfigStd.Marshal(data) + if err != nil { logger.Warn(fmt.Sprintf("Failed to encode JSON response: %v", err)) SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to encode response: %v", err)) + return } + ctx.SetBody(append(body, '\n')) } // SendError sends a BifrostError response diff --git a/transports/bifrost-http/handlers/utils_test.go b/transports/bifrost-http/handlers/utils_test.go index 583d49dd2c2..6547d41049a 100644 --- a/transports/bifrost-http/handlers/utils_test.go +++ b/transports/bifrost-http/handlers/utils_test.go @@ -1,79 +1,89 @@ package handlers import ( - "errors" - "net/http" - "net/http/httptest" - "os" - "path/filepath" + "strings" "testing" + + "github.com/valyala/fasthttp" ) -// TestIsUniqueConstraintError recognizes common database unique-constraint messages. -func TestIsUniqueConstraintError(t *testing.T) { - cases := []string{ - "UNIQUE constraint failed: enterprise_access_profiles.name", - `pq: duplicate key value violates unique constraint "idx_access_profiles_name"`, - "Error 1062: Duplicate entry 'profile-a' for key 'enterprise_access_profiles.name'", - } - for _, tc := range cases { - if !IsUniqueConstraintError(errors.New(tc)) { - t.Fatalf("IsUniqueConstraintError(%q)=false, want true", tc) - } - } - if IsUniqueConstraintError(errors.New("connection refused")) { - t.Fatalf("non-unique error should not match") +// TestSendJSON_StandardBytes pins the exact json.NewEncoder byte contract that the +// switch to sonic must not regress: HTML escaping, sorted map keys (deterministic +// output), and a trailing newline. The input has out-of-order keys and HTML +// metacharacters so all three properties show up in one exact-bytes comparison. +// sonic.ConfigDefault would break every one of them. +func TestSendJSON_StandardBytes(t *testing.T) { + ctx := &fasthttp.RequestCtx{} + SendJSON(ctx, map[string]string{"z": "1", "a": "&"}) + + got := string(ctx.Response.Body()) + want := `{"a":"\u003cb\u003e\u0026\u003c/b\u003e","z":"1"}` + "\n" + if got != want { + t.Errorf("SendJSON bytes mismatch:\n got %q\n want %q", got, want) } } -// TestIsUniqueConstraintError_Identifiers narrows matches to requested fields or indexes. -func TestIsUniqueConstraintError_Identifiers(t *testing.T) { - err := errors.New(`pq: duplicate key value violates unique constraint "idx_access_profiles_name"`) - if !IsUniqueConstraintError(err, "idx_access_profiles_name") { - t.Fatalf("identifier match returned false") - } - if IsUniqueConstraintError(err, "enterprise_users.email") { - t.Fatalf("unrelated identifier matched") +// TestSendJSON_Deterministic pins that a many-key map serializes byte-identically +// every time. Without key sorting (ConfigDefault) Go's randomized map iteration +// would make this flaky — the regression we must never reintroduce. +func TestSendJSON_Deterministic(t *testing.T) { + m := map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5, "f": 6, "g": 7, "h": 8} + first := &fasthttp.RequestCtx{} + SendJSON(first, m) + want := string(first.Response.Body()) + for i := 0; i < 50; i++ { + ctx := &fasthttp.RequestCtx{} + SendJSON(ctx, m) + if got := string(ctx.Response.Body()); got != want { + t.Fatalf("non-deterministic output on iter %d:\n want %s\n got %s", i, want, got) + } } } -func TestCheckURLAccessibility_FileExists(t *testing.T) { - f, err := os.CreateTemp(t.TempDir(), "pricing-*.json") - if err != nil { - t.Fatal(err) - } - f.Close() +// TestSendJSONWithStatus_StandardBytes mirrors SendJSON's exact-bytes contract and +// pins the custom status code. +func TestSendJSONWithStatus_StandardBytes(t *testing.T) { + ctx := &fasthttp.RequestCtx{} + SendJSONWithStatus(ctx, map[string]string{"z": "1", "a": ""}, fasthttp.StatusCreated) - if err := checkURLAccessibility("file://" + f.Name()); err != nil { - t.Fatalf("expected no error for existing file, got: %v", err) + if got := ctx.Response.StatusCode(); got != fasthttp.StatusCreated { + t.Errorf("status = %d, want %d", got, fasthttp.StatusCreated) } -} - -func TestCheckURLAccessibility_FileMissing(t *testing.T) { - path := filepath.Join(t.TempDir(), "nonexistent.json") - if err := checkURLAccessibility("file://" + path); err == nil { - t.Fatal("expected error for missing file, got nil") + got := string(ctx.Response.Body()) + want := `{"a":"\u003cx\u003e","z":"1"}` + "\n" + if got != want { + t.Errorf("SendJSONWithStatus bytes mismatch:\n got %q\n want %q", got, want) } } -func TestCheckURLAccessibility_HTTP200(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) - })) - defer srv.Close() +// TestSendJSON_MarshalError pins the marshal-failure early-return: an unsupported +// value (a channel) must yield HTTP 500 and the SendError body, never a partial or +// panicking write. Guards the err != nil branch that the happy-path tests skip. +func TestSendJSON_MarshalError(t *testing.T) { + SetLogger(&mockLogger{}) // error path logs a warning; shared no-op logger + ctx := &fasthttp.RequestCtx{} + SendJSON(ctx, make(chan int)) // channels are unmarshalable - if err := checkURLAccessibility(srv.URL); err != nil { - t.Fatalf("expected no error for HTTP 200, got: %v", err) + if got := ctx.Response.StatusCode(); got != fasthttp.StatusInternalServerError { + t.Errorf("status = %d, want %d", got, fasthttp.StatusInternalServerError) + } + if body := string(ctx.Response.Body()); !strings.Contains(body, "Failed to encode response") { + t.Errorf("expected SendError body, got %q", body) } } -func TestCheckURLAccessibility_HTTPNon200(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusNotFound) - })) - defer srv.Close() +// TestSendJSONWithStatus_MarshalError mirrors the marshal-failure early-return for +// the custom-status helper: the 500 from SendError must win over the requested +// status code. +func TestSendJSONWithStatus_MarshalError(t *testing.T) { + SetLogger(&mockLogger{}) + ctx := &fasthttp.RequestCtx{} + SendJSONWithStatus(ctx, make(chan int), fasthttp.StatusCreated) - if err := checkURLAccessibility(srv.URL); err == nil { - t.Fatal("expected error for HTTP 404, got nil") + if got := ctx.Response.StatusCode(); got != fasthttp.StatusInternalServerError { + t.Errorf("status = %d, want %d (SendError must override the requested status)", got, fasthttp.StatusInternalServerError) + } + if body := string(ctx.Response.Body()); !strings.Contains(body, "Failed to encode response") { + t.Errorf("expected SendError body, got %q", body) } }