Repository: zeromicro/go-zero
Pull request #2314 (2022-08-28 07:39:17)

safemap add Range method
safemap add Range method .


Patch:
diff --git a/core/collection/safemap.go b/core/collection/safemap.go
index 58e36d4e..7be1e4af 100644
--- a/core/collection/safemap.go
+++ b/core/collection/safemap.go
@@ -94,3 +94,31 @@ func (m *SafeMap) Size() int {
 	m.lock.RUnlock()
 	return size
 }
+
+// Range calls f sequentially for each key and value present in the map.
+// If f returns false, range stops the iteration.
+// Range m .
+func (m *SafeMap) Range(f func(key, value interface{}) bool) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+	var wg sync.WaitGroup
+
+	wg.Add(2)
+	go func() {
+		for k, v := range m.dirtyOld {
+			if !f(k, v) {
+				break
+			}
+		}
+		wg.Done()
+	}()
+	go func() {
+		for k, v := range m.dirtyNew {
+			if !f(k, v) {
+				break
+			}
+		}
+		wg.Done()
+	}()
+	wg.Wait()
+}
diff --git a/rest/httpc/responses.go b/rest/httpc/responses.go
index 704d0a9d..7207cdc7 100644
--- a/rest/httpc/responses.go
+++ b/rest/httpc/responses.go
@@ -1,6 +1,8 @@
 package httpc
 
 import (
+	"bytes"
+	"io"
 	"net/http"
 	"strings"
 
@@ -27,13 +29,24 @@ func ParseHeaders(resp *http.Response, val interface{}) error {
 func ParseJsonBody(resp *http.Response, val interface{}) error {
 	defer resp.Body.Close()
 
-	if withJsonBody(resp) {
-		return mapping.UnmarshalJsonReader(resp.Body, val)
+	if isContentTypeJson(resp) {
+		if resp.ContentLength > 0 {
+			return mapping.UnmarshalJsonReader(resp.Body, val)
+		}
+
+		var buf bytes.Buffer
+		if _, err := io.Copy(&buf, resp.Body); err != nil {
+			return err
+		}
+
+		if buf.Len() > 0 {
+			return mapping.UnmarshalJsonReader(&buf, val)
+		}
 	}
 
 	return mapping.UnmarshalJsonMap(nil, val)
 }
 
-func withJsonBody(r *http.Response) bool {
-	return r.ContentLength > 0 && strings.Contains(r.Header.Get(header.ContentType), header.ApplicationJson)
+func isContentTypeJson(r *http.Response) bool {
+	return strings.Contains(r.Header.Get(header.ContentType), header.ApplicationJson)
 }


Test patch:
diff --git a/core/collection/safemap_test.go b/core/collection/safemap_test.go
index 46e64749..4e9ca1d7 100644
--- a/core/collection/safemap_test.go
+++ b/core/collection/safemap_test.go
@@ -1,6 +1,7 @@
 package collection
 
 import (
+	"go.uber.org/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -107,3 +108,42 @@ func testSafeMapWithParameters(t *testing.T, size, exception int) {
 		}
 	}
 }
+
+func TestSafeMap_Range(t *testing.T) {
+	const (
+		size       = 100000
+		exception1 = 5
+		exception2 = 500
+	)
+	m := NewSafeMap()
+	m_new := NewSafeMap()
+
+	for i := 0; i < size; i++ {
+		m.Set(i, i)
+	}
+	for i := 0; i < size; i++ {
+		if i%exception1 == 0 {
+			m.Del(i)
+		}
+	}
+
+	for i := size; i < size<<1; i++ {
+		m.Set(i, i)
+	}
+	for i := size; i < size<<1; i++ {
+		if i%exception2 != 0 {
+			m.Del(i)
+		}
+	}
+
+	count := atomic.Int32{}
+	m.Range(func(k, v interface{}) bool {
+		count.Add(1)
+		m_new.Set(k, v)
+		return true
+	})
+	assert.Equal(t, int(count.Load()), m.Size())
+	assert.Equal(t, m.dirtyNew, m_new.dirtyNew)
+	assert.Equal(t, m.dirtyOld, m_new.dirtyOld)
+
+}
diff --git a/rest/httpc/responses_test.go b/rest/httpc/responses_test.go
index 6f1a2d1e..31369ee1 100644
--- a/rest/httpc/responses_test.go
+++ b/rest/httpc/responses_test.go
@@ -1,6 +1,7 @@
 package httpc
 
 import (
+	"errors"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -83,3 +84,97 @@ func TestParseWithZeroValue(t *testing.T) {
 	assert.Equal(t, 0, val.Foo)
 	assert.Equal(t, 0, val.Bar)
 }
+
+func TestParseWithNegativeContentLength(t *testing.T) {
+	var val struct {
+		Bar int `json:"bar"`
+	}
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set(header.ContentType, header.JsonContentType)
+		w.Write([]byte(`{"bar":0}`))
+	}))
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+
+	tests := []struct {
+		name   string
+		length int64
+	}{
+		{
+			name:   "negative",
+			length: -1,
+		},
+		{
+			name:   "zero",
+			length: 0,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			resp, err := DoRequest(req)
+			resp.ContentLength = test.length
+			assert.Nil(t, err)
+			assert.Nil(t, Parse(resp, &val))
+			assert.Equal(t, 0, val.Bar)
+		})
+	}
+}
+
+func TestParseWithNegativeContentLengthNoBody(t *testing.T) {
+	var val struct{}
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set(header.ContentType, header.JsonContentType)
+	}))
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+
+	tests := []struct {
+		name   string
+		length int64
+	}{
+		{
+			name:   "negative",
+			length: -1,
+		},
+		{
+			name:   "zero",
+			length: 0,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			resp, err := DoRequest(req)
+			resp.ContentLength = test.length
+			assert.Nil(t, err)
+			assert.Nil(t, Parse(resp, &val))
+		})
+	}
+}
+
+func TestParseJsonBody_BodyError(t *testing.T) {
+	var val struct{}
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set(header.ContentType, header.JsonContentType)
+	}))
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+
+	resp, err := DoRequest(req)
+	resp.ContentLength = -1
+	resp.Body = mockedReader{}
+	assert.Nil(t, err)
+	assert.NotNil(t, Parse(resp, &val))
+}
+
+type mockedReader struct{}
+
+func (m mockedReader) Close() error {
+	return nil
+}
+
+func (m mockedReader) Read(p []byte) (n int, err error) {
+	return 0, errors.New("dummy")
+}