Skip to content

Commit

Permalink
fix: Avoid memory copy in RESTful search and query (#37674)
Browse files Browse the repository at this point in the history
Custom `jsonRender` that encodes JSON data directly into the response
stream, it uses less memory since it does not buffer the entire JSON
structure before sending it, unlike `c.JSON` in `HTTPReturn`, which
serializes the JSON fully in memory before writing it to the response.
Benchmark testing shows that, using the custom render incurs no
performance loss and reduces memory consumption by nearly 50%.

issue: #37671

---------

Signed-off-by: bigsheeper <[email protected]>
  • Loading branch information
bigsheeper authored Nov 19, 2024
1 parent 3bbbf3c commit b3fc530
Show file tree
Hide file tree
Showing 19 changed files with 437 additions and 9 deletions.
16 changes: 16 additions & 0 deletions internal/distributed/proxy/httpserver/constant.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
Expand Down
16 changes: 16 additions & 0 deletions internal/distributed/proxy/httpserver/handler.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
Expand Down
16 changes: 16 additions & 0 deletions internal/distributed/proxy/httpserver/handler_test.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
Expand Down
22 changes: 19 additions & 3 deletions internal/distributed/proxy/httpserver/handler_v1.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
Expand Down Expand Up @@ -538,7 +554,7 @@ func (h *HandlersV1) query(c *gin.Context) {
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
HTTPReturn(c, http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
}
}
}
Expand Down Expand Up @@ -610,7 +626,7 @@ func (h *HandlersV1) get(c *gin.Context) {
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
HTTPReturn(c, http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
}
}
}
Expand Down Expand Up @@ -956,7 +972,7 @@ func (h *HandlersV1) search(c *gin.Context) {
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
HTTPReturn(c, http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
}
}
}
Expand Down
20 changes: 18 additions & 2 deletions internal/distributed/proxy/httpserver/handler_v1_test.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
Expand Down Expand Up @@ -514,7 +530,7 @@ func TestQuery(t *testing.T) {
name: "query success",
mp: mp5,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":[{\"book_id\":1,\"book_intro\":[0.1,0.11],\"word_count\":1000},{\"book_id\":2,\"book_intro\":[0.2,0.22],\"word_count\":2000},{\"book_id\":3,\"book_intro\":[0.3,0.33],\"word_count\":3000}]}",
expectedBody: "{\"code\":200,\"data\":[{\"book_id\":1,\"book_intro\":[0.1,0.11],\"word_count\":1000},{\"book_id\":2,\"book_intro\":[0.2,0.22],\"word_count\":2000},{\"book_id\":3,\"book_intro\":[0.3,0.33],\"word_count\":3000}]}\n",
})

for _, tt := range testCases {
Expand Down Expand Up @@ -1263,7 +1279,7 @@ func TestSearch(t *testing.T) {
name: "search success",
mp: mp5,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":[{\"book_id\":1,\"book_intro\":[0.1,0.11],\"distance\":0.01,\"word_count\":1000},{\"book_id\":2,\"book_intro\":[0.2,0.22],\"distance\":0.04,\"word_count\":2000},{\"book_id\":3,\"book_intro\":[0.3,0.33],\"distance\":0.09,\"word_count\":3000}]}",
expectedBody: "{\"code\":200,\"data\":[{\"book_id\":1,\"book_intro\":[0.1,0.11],\"distance\":0.01,\"word_count\":1000},{\"book_id\":2,\"book_intro\":[0.2,0.22],\"distance\":0.04,\"word_count\":2000},{\"book_id\":3,\"book_intro\":[0.3,0.33],\"distance\":0.09,\"word_count\":3000}]}\n",
})

for _, tt := range testCases {
Expand Down
24 changes: 20 additions & 4 deletions internal/distributed/proxy/httpserver/handler_v2.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
Expand Down Expand Up @@ -660,7 +676,7 @@ func (h *HandlersV2) query(ctx context.Context, c *gin.Context, anyReq any, dbNa
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
HTTPReturn(c, http.StatusOK, gin.H{
HTTPReturnStream(c, http.StatusOK, gin.H{
HTTPReturnCode: merr.Code(nil),
HTTPReturnData: outputData,
HTTPReturnCost: proxy.GetCostValue(queryResp.GetStatus()),
Expand Down Expand Up @@ -708,7 +724,7 @@ func (h *HandlersV2) get(ctx context.Context, c *gin.Context, anyReq any, dbName
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
HTTPReturn(c, http.StatusOK, gin.H{
HTTPReturnStream(c, http.StatusOK, gin.H{
HTTPReturnCode: merr.Code(nil),
HTTPReturnData: outputData,
HTTPReturnCost: proxy.GetCostValue(queryResp.GetStatus()),
Expand Down Expand Up @@ -1037,7 +1053,7 @@ func (h *HandlersV2) search(ctx context.Context, c *gin.Context, anyReq any, dbN
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
HTTPReturn(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost})
HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost})
}
}
}
Expand Down Expand Up @@ -1129,7 +1145,7 @@ func (h *HandlersV2) advancedSearch(ctx context.Context, c *gin.Context, anyReq
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
HTTPReturn(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost})
HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost})
}
}
}
Expand Down
16 changes: 16 additions & 0 deletions internal/distributed/proxy/httpserver/handler_v2_test.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
Expand Down
44 changes: 44 additions & 0 deletions internal/distributed/proxy/httpserver/json_render.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
"net/http"

"github.com/milvus-io/milvus/internal/json"
)

var jsonContentType = []string{"application/json; charset=utf-8"}

type jsonRender struct {
Data any
}

// Render writes data with custom ContentType.
func (r jsonRender) Render(w http.ResponseWriter) error {
r.WriteContentType(w)
encoder := json.NewEncoder(w)
return encoder.Encode(r.Data)
}

// WriteContentType writes JSON ContentType.
func (r jsonRender) WriteContentType(w http.ResponseWriter) {
header := w.Header()
if val := header["Content-Type"]; len(val) == 0 {
header["Content-Type"] = jsonContentType
}
}
101 changes: 101 additions & 0 deletions internal/distributed/proxy/httpserver/json_render_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

import (
"math/rand"
"net/http"
"net/http/httptest"
"strings"
"testing"

"github.com/gin-gonic/gin"
)

var testResult = gin.H{
"message": "This is a test message",
"data": make(map[string]interface{}),
}

func init() {
const chunkSize = 1024
rs := randomString(chunkSize)
var sb strings.Builder
for sb.Len() < 10*1024*1024 {
sb.WriteString(rs)
}
testResult["data"] = sb.String()
}

func randomString(length int) string {
const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
var result strings.Builder
for i := 0; i < length; i++ {
result.WriteByte(charset[rand.Intn(len(charset))])
}
return result.String()
}

func BenchmarkHTTPReturn(b *testing.B) {
// Set Gin to test mode to prevent output to stdout
gin.SetMode(gin.TestMode)

b.Run("test HTTPReturn", func(b *testing.B) {
router := gin.New()
router.GET("/test1", func(c *gin.Context) {
HTTPReturn(c, http.StatusOK, testResult)
})

b.ResetTimer()
for i := 0; i < b.N; i++ {
req := httptest.NewRequest("GET", "/test1", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
}

b.ReportAllocs()
})

b.Run("test HTTPReturnStream", func(b *testing.B) {
router := gin.New()
router.GET("/test2", func(c *gin.Context) {
HTTPReturnStream(c, http.StatusOK, testResult)
})

b.ResetTimer()
for i := 0; i < b.N; i++ {
req := httptest.NewRequest("GET", "/test2", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
}

b.ReportAllocs()
})
}

// goos: linux
// goarch: amd64
// pkg: github.com/milvus-io/milvus/internal/distributed/proxy/httpserver
// cpu: Intel(R) Core(TM) i7-8700 CPU @ 3.20GHz
// BenchmarkHTTPReturn
// BenchmarkHTTPReturn/test_HTTPReturn
// BenchmarkHTTPReturn/test_HTTPReturn-12 87 13127452 ns/op 27992718 B/op 34 allocs/op
// BenchmarkHTTPReturn/test_HTTPReturnStream
// BenchmarkHTTPReturn/test_HTTPReturnStream-12 87 12804875 ns/op 14361636 B/op 31 allocs/op
// PASS
//
// Process finished with the exit code 0
16 changes: 16 additions & 0 deletions internal/distributed/proxy/httpserver/request.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,19 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package httpserver

type CreateCollectionReq struct {
Expand Down
Loading

0 comments on commit b3fc530

Please sign in to comment.