don't use io.read-all (micro-optimization)

* add cos.ReadAll and cos.ReadAllN * with minor refactoring (htrun) Signed-off-by: Alex Aizman <[email protected]>
NVIDIA · Jul 21, 2024 · 6c8b7c0 · 6c8b7c0
1 parent 4cfb1f2
commit 6c8b7c0
Show file tree

Hide file tree

Showing 32 changed files with 433 additions and 399 deletions.
diff --git a/ais/backend/gcp.go b/ais/backend/gcp.go
@@ -462,7 +462,7 @@ func readCredFile() (projectID string) {
 	if err != nil {
 		return
 	}
-	b, err := io.ReadAll(credFile)
+	b, err := cos.ReadAll(credFile)
 	credFile.Close()
 	if err != nil {
 		return

diff --git a/ais/htcommon.go b/ais/htcommon.go
@@ -412,8 +412,13 @@ var (
 	_ cresv = cresBsumm{}
 )
 
-func (res *callResult) read(body io.Reader)  { res.bytes, res.err = io.ReadAll(body) }
-func (res *callResult) jread(body io.Reader) { res.err = jsoniter.NewDecoder(body).Decode(res.v) }
+func (res *callResult) read(body io.Reader, size int64) {
+	res.bytes, res.err = cos.ReadAllN(body, size)
+}
+
+func (res *callResult) jread(body io.Reader) {
+	res.err = jsoniter.NewDecoder(body).Decode(res.v)
+}
 
 func (res *callResult) mread(body io.Reader) {
 	vv, ok := res.v.(msgp.Decodable)

diff --git a/ais/htrun.go b/ais/htrun.go
@@ -647,7 +647,7 @@ func (h *htrun) call(args *callArgs, smap *smapX) (res *callResult) {
 	if res.err != nil {
 		res.details = fmt.Sprintf("FATAL: failed to create HTTP request %s %s: %v",
 			args.req.Method, args.req.URL(), res.err)
-		return
+		return res
 	}
 
 	req.Header.Set(apc.HdrCallerID, h.SID())
@@ -663,9 +663,19 @@ func (h *htrun) call(args *callArgs, smap *smapX) (res *callResult) {
 	resp, res.err = client.Do(req)
 	if res.err != nil {
 		res.details = dfltDetail // tcp level, e.g.: connection refused
-		return
+		return res
+	}
+
+	_doResp(args, req, resp, res)
+	resp.Body.Close()
+
+	if sid != unknownDaemonID {
+		h.keepalive.heardFrom(sid)
 	}
-	defer resp.Body.Close()
+	return res
+}
+
+func _doResp(args *callArgs, req *http.Request, resp *http.Response, res *callResult) {
 	res.status = resp.StatusCode
 	res.header = resp.Header
 
@@ -684,25 +694,14 @@ func (h *htrun) call(args *callArgs, smap *smapX) (res *callResult) {
 		return
 	}
 
-	// read and decode via call result value (`cresv`), if provided
+	// read and decode via call-result-value (`cresv`), if provided;
 	// othwerwise, read and return bytes for the caller to unmarshal
 	if args.cresv != nil {
 		res.v = args.cresv.newV()
 		args.cresv.read(res, resp.Body)
-		if res.err != nil {
-			return
-		}
 	} else {
-		res.read(resp.Body)
-		if res.err != nil {
-			return
-		}
+		res.read(resp.Body, resp.ContentLength)
 	}
-
-	if sid != unknownDaemonID {
-		h.keepalive.heardFrom(sid)
-	}
-	return
 }
 
 //

diff --git a/ais/proxy.go b/ais/proxy.go
@@ -7,7 +7,6 @@ package ais
 import (
 	"errors"
 	"fmt"
-	"io"
 	"net"
 	"net/http"
 	"net/url"
@@ -3007,7 +3006,7 @@ func (p *proxy) dsortHandler(w http.ResponseWriter, r *http.Request) {
 	case http.MethodPost:
 		// - validate request, check input_bck and output_bck
 		// - start dsort
-		body, err := io.ReadAll(r.Body)
+		body, err := cos.ReadAllN(r.Body, r.ContentLength)
 		if err != nil {
 			p.writeErrStatusf(w, r, http.StatusInternalServerError, "failed to receive dsort request: %v", err)
 			return

diff --git a/ais/prxdl.go b/ais/prxdl.go
@@ -6,7 +6,6 @@ package ais
 
 import (
 	"fmt"
-	"io"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -87,7 +86,7 @@ func (p *proxy) httpdlpost(w http.ResponseWriter, r *http.Request) {
 
 	jobID := dload.PrefixJobID + cos.GenUUID() // prefix to visually differentiate vs. xaction IDs
 
-	body, err := io.ReadAll(r.Body)
+	body, err := cos.ReadAllN(r.Body, r.ContentLength)
 	if err != nil {
 		p.writeErrStatusf(w, r, http.StatusInternalServerError, "failed to receive download request: %v", err)
 		return

diff --git a/ais/prxetl.go b/ais/prxetl.go
@@ -5,7 +5,6 @@
 package ais
 
 import (
-	"io"
 	"net/http"
 	"net/url"
 	"reflect"
@@ -100,7 +99,7 @@ func (p *proxy) handleETLPut(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	b, err := io.ReadAll(r.Body)
+	b, err := cos.ReadAll(r.Body)
 	if err != nil {
 		p.writeErr(w, r, err)
 		return

diff --git a/ais/prxnotif_internal_test.go b/ais/prxnotif_internal_test.go
@@ -6,7 +6,6 @@ package ais
 
 import (
 	"bytes"
-	"io"
 	"net/http"
 	"net/http/httptest"
 	"time"
@@ -138,7 +137,7 @@ var _ = Describe("Notifications xaction test", func() {
 			writer := httptest.NewRecorder()
 			n.handler(writer, req)
 			resp := writer.Result()
-			respBody, _ := io.ReadAll(resp.Body)
+			respBody, _ := cos.ReadAllN(resp.Body, resp.ContentLength)
 			resp.Body.Close()
 			Expect(resp.StatusCode).To(BeEquivalentTo(expectedStatus))
 			return respBody

diff --git a/ais/s3/presigned.go b/ais/s3/presigned.go
@@ -70,7 +70,7 @@ func (pts *PresignedReq) Do(client *http.Client) (*PresignedResp, error) {
 	}
 	defer resp.BodyR.Close()
 
-	output, err := io.ReadAll(resp.BodyR)
+	output, err := cos.ReadAll(resp.BodyR)
 	if err != nil {
 		return &PresignedResp{StatusCode: http.StatusBadRequest}, fmt.Errorf("failed to read response body: %v", err)
 	}
@@ -111,7 +111,7 @@ func (pts *PresignedReq) DoReader(client *http.Client) (*PresignedResp, error) {
 	}
 
 	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
-		output, _ := io.ReadAll(resp.Body)
+		output, _ := cos.ReadAll(resp.Body)
 		resp.Body.Close()
 		return &PresignedResp{StatusCode: resp.StatusCode}, fmt.Errorf("invalid status: %d, output: %s", resp.StatusCode, string(output))
 	}

diff --git a/ais/tgtetl.go b/ais/tgtetl.go
@@ -6,7 +6,6 @@ package ais
 
 import (
 	"fmt"
-	"io"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -56,7 +55,7 @@ func (t *target) handleETLPut(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	b, err := io.ReadAll(r.Body)
+	b, err := cos.ReadAll(r.Body)
 	if err != nil {
 		t.writeErr(w, r, err)
 		return

diff --git a/ais/tgts3mpt.go b/ais/tgts3mpt.go
@@ -215,7 +215,7 @@ func (t *target) completeMpt(w http.ResponseWriter, r *http.Request, items []str
 		return
 	}
 
-	output, err := io.ReadAll(r.Body)
+	output, err := cos.ReadAllN(r.Body, r.ContentLength)
 	if err != nil {
 		s3.WriteErr(w, r, err, http.StatusBadRequest)
 		return

diff --git a/api/client.go b/api/client.go
@@ -279,7 +279,7 @@ func (reqParams *ReqParams) readStr(resp *http.Response, out *string) error {
 	if err := reqParams.checkResp(resp); err != nil {
 		return err
 	}
-	b, err := io.ReadAll(resp.Body)
+	b, err := cos.ReadAllN(resp.Body, resp.ContentLength)
 	if err != nil {
 		return fmt.Errorf("failed to read response: %w", err)
 	}
@@ -351,7 +351,7 @@ func (reqParams *ReqParams) checkResp(resp *http.Response) error {
 		}
 	}
 
-	b, _ := io.ReadAll(resp.Body)
+	b, _ := cos.ReadAllN(resp.Body, resp.ContentLength)
 	if len(b) == 0 {
 		if resp.StatusCode == http.StatusServiceUnavailable {
 			msg := fmt.Sprintf("[%s]: starting up, please try again later...", http.StatusText(http.StatusServiceUnavailable))

diff --git a/api/etl.go b/api/etl.go
@@ -62,7 +62,7 @@ func ETLGetInitMsg(params BaseParams, etlName string) (etl.InitMsg, error) {
 	}
 	defer cos.Close(r)
 
-	b, err := io.ReadAll(r)
+	b, err := cos.ReadAll(r)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response: %w", err)
 	}

diff --git a/api/stats.go b/api/stats.go
@@ -5,7 +5,6 @@
 package api
 
 import (
-	"io"
 	"net/http"
 	"net/url"
 
@@ -112,7 +111,7 @@ func GetAnyStats(bp BaseParams, sid, what string) (out []byte, err error) {
 	if err != nil {
 		return nil, err
 	}
-	out, err = io.ReadAll(resp.Body)
+	out, err = cos.ReadAllN(resp.Body, resp.ContentLength)
 	cos.DrainReader(resp.Body)
 	resp.Body.Close()
 	FreeRp(reqParams)

diff --git a/bench/tools/aisloader/client.go b/bench/tools/aisloader/client.go
@@ -557,7 +557,7 @@ func readDiscard(r *http.Response, tag, cksumType string) (int64, string, error)
 		cksumValue string
 	)
 	if r.StatusCode >= http.StatusBadRequest {
-		bytes, err := io.ReadAll(r.Body)
+		bytes, err := cos.ReadAll(r.Body)
 		if err == nil {
 			return 0, "", fmt.Errorf("bad status %d from %s, response: %s", r.StatusCode, tag, string(bytes))
 		}

diff --git a/bench/tools/aisloader/run.go b/bench/tools/aisloader/run.go
@@ -27,7 +27,6 @@ import (
 	"errors"
 	"flag"
 	"fmt"
-	"io"
 	"math"
 	"math/rand/v2"
 	"os"
@@ -816,7 +815,7 @@ func _init(p *params) (err error) {
 		if err != nil {
 			return err
 		}
-		etlSpec, err := io.ReadAll(fh)
+		etlSpec, err := cos.ReadAll(fh)
 		fh.Close()
 		if err != nil {
 			return err

diff --git a/cmd/authn/aisreq.go b/cmd/authn/aisreq.go
@@ -7,7 +7,6 @@ package main
 import (
 	"bytes"
 	"fmt"
-	"io"
 	"net/http"
 	"sync"
 	"time"
@@ -125,7 +124,7 @@ func (m *mgr) call(method, proxyURL, path string, injson []byte, tag string) err
 		resp, err := client.Do(req)
 		if resp != nil {
 			if resp.Body != nil {
-				msg, _ = io.ReadAll(resp.Body)
+				msg, _ = cos.ReadAll(resp.Body)
 				resp.Body.Close()
 			}
 		}