Make Write message type more flexible

Signed-off-by: Saswata Mukherjee <[email protected]>

Author: saswatamcode

api/prometheus/v1/remote/remote_api.go

@@ -134,30 +134,58 @@ type vtProtoEnabled interface {
 	MarshalToSizedBufferVT(dAtA []byte) (int, error)
 }
 
+type gogoProtoEnabled interface {
+	Size() (n int)
+	MarshalToSizedBuffer(dAtA []byte) (n int, err error)
+}
+
+// Sort of a hack to identify v2 requests.
+// Under any marshaling scheme, v2 requests have a `Symbols` field of type []string.
+// So would always have a `GetSymbols()` method which doesn't rely on any other types.
+type v2Request interface {
+	GetSymbols() []string
+}
+
 // Write writes given, non-empty, protobuf message to a remote storage.
 // The https://github.com/planetscale/vtprotobuf methods will be used if your msg
-// supports those (e.g. SizeVT() and MarshalToSizedBufferVT(...)), for efficiency.
-func (r *API) Write(ctx context.Context, msg proto.Message) (_ WriteResponseStats, err error) {
+// supports those (e.g. SizeVT() and MarshalToSizedBufferVT(...)), for efficiency
+// or https://github.com/gogo/protobuf methods (e.g. Size() and MarshalToSizedBuffer(...))
+// will be used if your msg supports those.
+func (r *API) Write(ctx context.Context, msg any) (_ WriteResponseStats, err error) {
 	// Detect content-type.
-	cType := WriteProtoFullName(proto.MessageName(msg))
+	cType := WriteProtoFullNameV1
+	if _, ok := msg.(v2Request); ok {
+		cType = WriteProtoFullNameV2
+	}
+
 	if err := cType.Validate(); err != nil {
 		return WriteResponseStats{}, err
 	}
 
 	// Encode the payload.
-	if emsg, ok := msg.(vtProtoEnabled); ok {
+	switch m := msg.(type) {
+	case vtProtoEnabled:
 		// Use optimized vtprotobuf if supported.
-		size := emsg.SizeVT()
+		size := m.SizeVT()
 		if len(r.reqBuf) < size {
 			r.reqBuf = make([]byte, size)
 		}
-		if _, err := emsg.MarshalToSizedBufferVT(r.reqBuf[:size]); err != nil {
+		if _, err := m.MarshalToSizedBufferVT(r.reqBuf[:size]); err != nil {
 			return WriteResponseStats{}, fmt.Errorf("encoding request %w", err)
 		}
-	} else {
+	case gogoProtoEnabled:
+		// Gogo proto if supported.
+		size := m.Size()
+		if len(r.reqBuf) < size {
+			r.reqBuf = make([]byte, size)
+		}
+		if _, err := m.MarshalToSizedBuffer(r.reqBuf[:size]); err != nil {
+			return WriteResponseStats{}, fmt.Errorf("encoding request %w", err)
+		}
+	case proto.Message:
 		// Generic proto.
 		r.reqBuf = r.reqBuf[:0]
-		r.reqBuf, err = (proto.MarshalOptions{}).MarshalAppend(r.reqBuf, msg)
+		r.reqBuf, err = (proto.MarshalOptions{}).MarshalAppend(r.reqBuf, m)
 		if err != nil {
 			return WriteResponseStats{}, fmt.Errorf("encoding request %w", err)
 		}

api/prometheus/v1/remote/remote_api_test.go

@@ -27,6 +27,7 @@ import (
 	"google.golang.org/protobuf/testing/protocmp"
 
 	"github.com/prometheus/client_golang/api"
+	writev1 "github.com/prometheus/client_golang/api/prometheus/v1/remote/genproto/v1"
 	writev2 "github.com/prometheus/client_golang/api/prometheus/v1/remote/genproto/v2"
 )
 
@@ -61,6 +62,7 @@ func TestRetryAfterDuration(t *testing.T) {
 
 type mockStorage struct {
 	v2Reqs []*writev2.Request
+	v1Reqs []*writev1.WriteRequest
 	protos []WriteProtoFullName
 
 	mockCode *int
@@ -72,14 +74,23 @@ func (m *mockStorage) Store(_ context.Context, msgFullName WriteProtoFullName, s
 		return w, *m.mockCode, m.mockErr
 	}
 
-	// This test expects v2 only.
-	r := &writev2.Request{}
-	if err := proto.Unmarshal(serializedRequest, r); err != nil {
-		return WriteResponseStats{}, http.StatusInternalServerError, err
+	if msgFullName == WriteProtoFullNameV1 {
+		r := &writev1.WriteRequest{}
+		if err := proto.Unmarshal(serializedRequest, r); err != nil {
+			return WriteResponseStats{}, http.StatusInternalServerError, err
+		}
+		m.v1Reqs = append(m.v1Reqs, r)
+		m.protos = append(m.protos, msgFullName)
+		return WriteResponseStats{}, http.StatusOK, nil
+	} else {
+		r := &writev2.Request{}
+		if err := proto.Unmarshal(serializedRequest, r); err != nil {
+			return WriteResponseStats{}, http.StatusInternalServerError, err
+		}
+		m.v2Reqs = append(m.v2Reqs, r)
+		m.protos = append(m.protos, msgFullName)
+		return stats(r), http.StatusOK, nil
 	}
-	m.v2Reqs = append(m.v2Reqs, r)
-	m.protos = append(m.protos, msgFullName)
-	return stats(r), http.StatusOK, nil
 }
 
 func testV2() *writev2.Request {
@@ -112,6 +123,20 @@ func testV2() *writev2.Request {
 	}
 }
 
+func testV1() *writev1.WriteRequest {
+	return &writev1.WriteRequest{
+		Timeseries: []*writev1.TimeSeries{
+			{
+				Labels: []*writev1.Label{
+					{Name: "__name__", Value: "metric1"},
+					{Name: "foo", Value: "bar1"},
+				},
+				Samples: []*writev1.Sample{{Value: 1.1, Timestamp: 1214141}, {Value: 1.5, Timestamp: 1214180}},
+			},
+		},
+	}
+}
+
 func stats(req *writev2.Request) (s WriteResponseStats) {
 	s.Confirmed = true
 	for _, ts := range req.Timeseries {
@@ -145,11 +170,22 @@ func TestRemoteAPI_Write_WithHandler(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
 	if diff := cmp.Diff(stats(req), s); diff != "" {
 		t.Fatal("unexpected stats", diff)
 	}
+
+	req2 := testV1()
+	_, err = client.Write(context.Background(), req2)
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	if len(mStore.v2Reqs) != 1 {
-		t.Fatal("expected 1 request stored, got", mStore.v2Reqs)
+		t.Fatal("expected 1 v2 request stored, got", mStore.v2Reqs)
+	}
+	if len(mStore.v1Reqs) != 1 {
+		t.Fatal("expected 1 v1 request stored, got", mStore.v1Reqs)
 	}
 	if diff := cmp.Diff(req, mStore.v2Reqs[0], protocmp.Transform()); diff != "" {
 		t.Fatal("unexpected request received", diff)

api/prometheus/v1/remote/remote_headers.go

@@ -19,8 +19,6 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
-
-	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
 const (
@@ -43,7 +41,7 @@ const (
 // WriteProtoFullName represents the fully qualified name of the protobuf message
 // to use in Remote write 1.0 and 2.0 protocols.
 // See https://prometheus.io/docs/specs/remote_write_spec_2_0/#protocol.
-type WriteProtoFullName protoreflect.FullName
+type WriteProtoFullName string
 
 const (
 	// WriteProtoFullNameV1 represents the `prometheus.WriteRequest` protobuf