Repository: zeromicro/go-zero
Pull request #2755 (2023-01-05 12:57:29)

feat: add batch inserter

codecov-commenter: # [Codecov](https://codecov.io/gh/zeromicro/go-zero/pull/2755?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#2755](https://codecov.io/gh/zeromicro/go-zero/pull/2755?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (a8b3fc0) into [master](https://codecov.io/gh/zeromicro/go-zero/commit/69d355eb4ba722144fdbae21677d8a7ca2c600be?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (69d355e) will **increase** coverage by `0.00%`.
> The diff coverage is `81.25%`.


| [Impacted Files](https://codecov.io/gh/zeromicro/go-zero/pull/2755?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [core/stores/mon/bulkinserter.go](https://codecov.io/gh/zeromicro/go-zero/pull/2755?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvbW9uL2J1bGtpbnNlcnRlci5nbw==) | `78.18% <81.25%> (+1.25%)` | :arrow_up: |
| [core/stores/sqlx/sqlconn.go](https://codecov.io/gh/zeromicro/go-zero/pull/2755?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvc3FseC9zcWxjb25uLmdv) | `56.37% <0.00%> (-1.48%)` | :arrow_down: |
| [core/stores/sqlx/tx.go](https://codecov.io/gh/zeromicro/go-zero/pull/2755?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvc3FseC90eC5nbw==) | `23.30% <0.00%> (+3.88%)` | :arrow_up: |


Patch:
diff --git a/core/stores/mon/bulkinserter.go b/core/stores/mon/bulkinserter.go
index c3fdfa23..74cbb7d9 100644
--- a/core/stores/mon/bulkinserter.go
+++ b/core/stores/mon/bulkinserter.go
@@ -25,8 +25,23 @@ type (
 	}
 )
 
-// NewBulkInserter returns a BulkInserter.
+// Deprecated. Use NewBatchInserter instead.
 func NewBulkInserter(coll *mongo.Collection, interval ...time.Duration) *BulkInserter {
+	return newBulkInserter(coll, interval...)
+}
+
+// NewBatchInserter returns a BulkInserter.
+func NewBatchInserter(coll Collection, interval ...time.Duration) (*BulkInserter, error) {
+	cloneColl, err := coll.Clone()
+	if err != nil {
+		return nil, err
+	}
+
+	return newBulkInserter(cloneColl, interval...), nil
+}
+
+// newBulkInserter returns a BulkInserter.
+func newBulkInserter(coll *mongo.Collection, interval ...time.Duration) *BulkInserter {
 	inserter := &dbInserter{
 		collection: coll,
 	}
diff --git a/rest/handler/tracinghandler.go b/rest/handler/tracinghandler.go
index 43ad9d2c..d3151569 100644
--- a/rest/handler/tracinghandler.go
+++ b/rest/handler/tracinghandler.go
@@ -1,24 +1,18 @@
 package handler
 
 import (
-	"bufio"
-	"errors"
-	"net"
 	"net/http"
 	"sync"
 
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/trace"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/propagation"
 	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
-const traceKeyStatusCode = "http.status_code"
-
 var notTracingSpans sync.Map
 
 // DontTraceSpan disable tracing for the specified span name.
@@ -26,40 +20,6 @@ func DontTraceSpan(spanName string) {
 	notTracingSpans.Store(spanName, lang.Placeholder)
 }
 
-type traceResponseWriter struct {
-	w    http.ResponseWriter
-	code int
-}
-
-// Flush implements the http.Flusher interface.
-func (w *traceResponseWriter) Flush() {
-	if flusher, ok := w.w.(http.Flusher); ok {
-		flusher.Flush()
-	}
-}
-
-// Hijack implements the http.Hijacker interface.
-func (w *traceResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	if hijacked, ok := w.w.(http.Hijacker); ok {
-		return hijacked.Hijack()
-	}
-
-	return nil, nil, errors.New("server doesn't support hijacking")
-}
-
-func (w *traceResponseWriter) Header() http.Header {
-	return w.w.Header()
-}
-
-func (w *traceResponseWriter) Write(data []byte) (int, error) {
-	return w.w.Write(data)
-}
-
-func (w *traceResponseWriter) WriteHeader(statusCode int) {
-	w.w.WriteHeader(statusCode)
-	w.code = statusCode
-}
-
 // TracingHandler return a middleware that process the opentelemetry.
 func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
@@ -89,21 +49,12 @@ func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 
 			// convenient for tracking error messages
 			propagator.Inject(spanCtx, propagation.HeaderCarrier(w.Header()))
-			trw := &traceResponseWriter{
-				w:    w,
-				code: http.StatusOK,
-			}
+
+			trw := &response.WithCodeResponseWriter{Writer: w, Code: http.StatusOK}
 			next.ServeHTTP(trw, r.WithContext(spanCtx))
 
-			span.SetAttributes(attribute.KeyValue{
-				Key:   traceKeyStatusCode,
-				Value: attribute.IntValue(trw.code),
-			})
-			if trw.code >= http.StatusBadRequest {
-				span.SetStatus(codes.Error, http.StatusText(trw.code))
-			} else {
-				span.SetStatus(codes.Ok, http.StatusText(trw.code))
-			}
+			span.SetAttributes(semconv.HTTPAttributesFromHTTPStatusCode(trw.Code)...)
+			span.SetStatus(semconv.SpanStatusFromHTTPStatusCodeAndSpanKind(trw.Code, oteltrace.SpanKindServer))
 		})
 	}
 }
diff --git a/rest/httpc/requests.go b/rest/httpc/requests.go
index 4ddb15a8..3685d5c6 100644
--- a/rest/httpc/requests.go
+++ b/rest/httpc/requests.go
@@ -190,7 +190,7 @@ func request(r *http.Request, cli client) (*http.Response, error) {
 	}
 
 	span.SetAttributes(semconv.HTTPAttributesFromHTTPStatusCode(resp.StatusCode)...)
-	span.SetStatus(semconv.SpanStatusFromHTTPStatusCode(resp.StatusCode))
+	span.SetStatus(semconv.SpanStatusFromHTTPStatusCodeAndSpanKind(resp.StatusCode, oteltrace.SpanKindClient))
 
 	return resp, err
 }


Test patch:
diff --git a/core/stores/mon/bulkinserter_test.go b/core/stores/mon/bulkinserter_test.go
index 0d6aa214..f2889d35 100644
--- a/core/stores/mon/bulkinserter_test.go
+++ b/core/stores/mon/bulkinserter_test.go
@@ -25,3 +25,21 @@ func TestBulkInserter(t *testing.T) {
 		bulk.Flush()
 	})
 }
+
+func TestBatchInserter(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
+		bulk, err := NewBatchInserter(createModel(mt).Collection)
+		assert.Equal(t, err, nil)
+		bulk.SetResultHandler(func(result *mongo.InsertManyResult, err error) {
+			assert.Nil(t, err)
+			assert.Equal(t, 2, len(result.InsertedIDs))
+		})
+		bulk.Insert(bson.D{{Key: "foo", Value: "bar"}})
+		bulk.Insert(bson.D{{Key: "foo", Value: "baz"}})
+		bulk.Flush()
+	})
+}
diff --git a/rest/handler/tracinghandler_test.go b/rest/handler/tracinghandler_test.go
index d4aebe2a..79a94c4d 100644
--- a/rest/handler/tracinghandler_test.go
+++ b/rest/handler/tracinghandler_test.go
@@ -148,46 +148,3 @@ func TestTraceResponseWriter(t *testing.T) {
 		})
 	}
 }
-
-func TestTraceHandler_traceResponseWriter(t *testing.T) {
-	writer := &traceResponseWriter{
-		w: httptest.NewRecorder(),
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-
-	writer = &traceResponseWriter{
-		w: mockedHijackable{httptest.NewRecorder()},
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-
-	writer = &traceResponseWriter{
-		w: httptest.NewRecorder(),
-	}
-	writer.WriteHeader(http.StatusBadRequest)
-	assert.NotNil(t, writer.Header())
-
-	writer = &traceResponseWriter{
-		w: httptest.NewRecorder(),
-	}
-	assert.NotPanics(t, func() {
-		writer.Flush()
-	})
-
-	writer = &traceResponseWriter{
-		w: mockedFlusher{httptest.NewRecorder()},
-	}
-	assert.NotPanics(t, func() {
-		writer.Flush()
-	})
-}
-
-type mockedFlusher struct {
-	*httptest.ResponseRecorder
-}
-
-func (m mockedFlusher) Flush() {
-}