measure.go

// Copyright 2022-2023 Tigris Data, Inc.
//
// Licensed 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 middleware

import (
	"context"
	"strconv"

	middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
	"github.com/rs/zerolog/log"
	api "github.com/tigrisdata/tigris/api/server/v1"
	"github.com/tigrisdata/tigris/errors"
	"github.com/tigrisdata/tigris/server/config"
	"github.com/tigrisdata/tigris/server/defaults"
	"github.com/tigrisdata/tigris/server/metrics"
	"github.com/tigrisdata/tigris/server/request"
	"github.com/tigrisdata/tigris/util"
	ulog "github.com/tigrisdata/tigris/util/log"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"
)

type wrappedStream struct {
	*middleware.WrappedServerStream
	measurement           *metrics.Measurement
	reqStatus             *metrics.RequestStatus
	setServerTimingHeader bool
}

func getNoMeasurementMethods() []string {
	return []string{
		api.HealthMethodName,
	}
}

func measureMethod(fullMethod string) bool {
	for _, method := range getNoMeasurementMethods() {
		if method == fullMethod {
			return false
		}
	}
	return true
}

func logError(ctx context.Context, err error, reqType string) {
	httpCode := 999 // Integer value for the unknown response code
	userAgent := defaults.UnknownValue
	message := defaults.UnknownValue
	var tigrisErr *api.TigrisError
	if grpcMetaData, grpcMetaDataFound := metadata.FromIncomingContext(ctx); grpcMetaDataFound {
		if userAgents, userAgentFound := grpcMetaData["user-agent"]; userAgentFound {
			userAgent = userAgents[0]
		}
	}
	if errors.As(err, &tigrisErr) {
		message = tigrisErr.Message
		httpCode = api.ToHTTPCode(tigrisErr.Code)
	}
	log.Error().Err(err).Int("HTTP response", httpCode).Str("message", message).Str("user-agent", userAgent).Str("request type", reqType).Msg("grpc request error")
}

func measureUnary() func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
		var reqStatus *metrics.RequestStatus
		if !measureMethod(info.FullMethod) {
			resp, err := handler(ctx, req)
			return resp, err
		}
		reqMetadata, err := request.GetRequestMetadataFromContext(ctx)
		ulog.E(err)
		reqStatus = metrics.NewRequestStatus(reqMetadata.GetNamespace())
		tags := reqMetadata.GetInitialTags()
		measurement := metrics.NewMeasurement(util.Service, info.FullMethod, metrics.GrpcSpanType, tags)
		measurement.AddTags(metrics.GetProjectBranchCollTags(reqMetadata.GetProject(), reqMetadata.GetBranch(), reqMetadata.GetCollection()))
		measurement.AddTags(map[string]string{
			"sub":   reqMetadata.Sub,
			"human": strconv.FormatBool(reqMetadata.IsHuman),
		})
		ctx = measurement.StartTracing(ctx, false)
		ctx = reqStatus.SaveRequestStatusToContext(ctx)
		resp, err := handler(ctx, req)
		if err != nil {
			// Request had an error
			measurement.SetError()
			measurement.CountErrorForScope(metrics.RequestsErrorCount, measurement.GetRequestErrorTags(err))
			measurement.AddReceivedBytes(proto.Size(req.(proto.Message)))
			_ = measurement.FinishWithError(ctx, err)
			measurement.RecordDuration(metrics.RequestsErrorRespTime, measurement.GetRequestErrorTags(err))
			logError(ctx, err, "unary")
			return nil, err
		}
		if err = setHeadersToUnary(ctx, resp, measurement); err != nil {
			return nil, err
		}

		// Request was ok
		measurement.SetNDocs(reqStatus.GetResultDocs())
		measurement.CountOkForScope(metrics.RequestsOkCount, measurement.GetRequestOkTags())
		measurement.AddReceivedBytes(proto.Size(req.(proto.Message)))
		measurement.AddSentBytes(proto.Size(resp.(proto.Message)))
		measurement.CountReceivedBytes(metrics.BytesReceived, measurement.GetNetworkTags())
		measurement.CountSentBytes(metrics.BytesSent, measurement.GetNetworkTags())
		_ = measurement.FinishTracing(ctx)
		measurement.RecordDuration(metrics.RequestsRespTime, measurement.GetRequestOkTags())
		// Global status and metrics related to them, config switches are handled inside these
		metrics.GlobalSt.RecordRequestToActiveChunk(reqStatus, reqMetadata.GetNamespace())
		measurement.CountUnits(reqStatus, measurement.GetGlobalStatusTags())
		return resp, err
	}
}

func measureStream() grpc.StreamServerInterceptor {
	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		var reqStatus *metrics.RequestStatus
		wrapped := &wrappedStream{WrappedServerStream: middleware.WrapServerStream(stream)}
		wrapped.WrappedContext = stream.Context()
		if !measureMethod(info.FullMethod) {
			err := handler(srv, wrapped)
			return err
		}
		reqMetadata, err := request.GetRequestMetadataFromContext(wrapped.WrappedContext)
		if err != nil {
			ulog.E(err)
		}
		reqStatus = metrics.NewRequestStatus(reqMetadata.GetNamespace())
		tags := reqMetadata.GetInitialTags()
		measurement := metrics.NewMeasurement(util.Service, info.FullMethod, metrics.GrpcSpanType, tags)
		wrapped.measurement = measurement
		wrapped.reqStatus = reqStatus
		wrapped.WrappedContext = reqStatus.SaveRequestStatusToContext(wrapped.WrappedContext)
		wrapped.WrappedContext = measurement.StartTracing(wrapped.WrappedContext, false)
		err = handler(srv, wrapped)
		if err != nil {
			measurement.SetError()
			measurement.CountErrorForScope(metrics.RequestsErrorCount, measurement.GetRequestErrorTags(err))
			_ = measurement.FinishWithError(wrapped.WrappedContext, err)
			measurement.RecordDuration(metrics.RequestsErrorRespTime, measurement.GetRequestErrorTags(err))
			logError(wrapped.WrappedContext, err, "stream")
			ulog.E(err)
		} else {
			ulog.E(setHeadersToStream(wrapped))
			measurement.CountOkForScope(metrics.RequestsOkCount, measurement.GetRequestOkTags())
			_ = measurement.FinishTracing(wrapped.WrappedContext)
			measurement.RecordDuration(metrics.RequestsRespTime, measurement.GetRequestOkTags())
		}
		// Global status and metrics related to them, config switches are handled inside these
		metrics.GlobalSt.RecordRequestToActiveChunk(reqStatus, reqMetadata.GetNamespace())
		measurement.CountUnits(reqStatus, measurement.GetGlobalStatusTags())
		// The actual bytes are counted in RecvMsg during streaming
		measurement.CountReceivedBytes(metrics.BytesReceived, measurement.GetNetworkTags())
		// The actual bytes are counted in SendMsg during streaming
		measurement.CountSentBytes(metrics.BytesSent, measurement.GetNetworkTags())
		wrapped.WrappedContext = reqStatus.SaveRequestStatusToContext(wrapped.WrappedContext)
		return err
	}
}

func (w *wrappedStream) RecvMsg(m any) error {
	recvErr := w.ServerStream.RecvMsg(m)
	if recvErr != nil {
		log.Error().Err(recvErr).Interface("message", m).Msg("stream receive message error")
	}
	if w.measurement == nil {
		return recvErr
	}

	if len(w.measurement.GetProjectCollTags()) == 0 {
		// The request is not tagged yet with db and collection, need to do it on the first message
		project, branch, coll := request.GetProjectAndBranchAndColl(m)
		reqMetadata, err := request.GetRequestMetadataFromContext(w.WrappedContext)
		if err != nil {
			if config.DefaultConfig.Metrics.DebugMessages {
				log.Debug().Str("error", err.Error()).Msg("error while getting request metadata, not measuring")
			}
			return recvErr
		}
		reqMetadata.SetProject(project)
		reqMetadata.SetBranch(branch)
		reqMetadata.SetCollection(coll)
		w.measurement.AddTags(map[string]string{
			"human": strconv.FormatBool(reqMetadata.IsHuman),
		})
		w.WrappedContext = reqMetadata.SaveToContext(w.WrappedContext)
		w.measurement.AddProjectBranchCollTags(project, branch, coll)
	}
	w.measurement.AddReceivedBytes(proto.Size(m.(proto.Message)))
	return recvErr
}

func (w *wrappedStream) SendMsg(m any) error {
	err := setHeadersToStream(w)
	if ulog.E(err) {
		return err
	}

	err = w.ServerStream.SendMsg(m)
	if err != nil {
		log.Error().Err(err).Interface("message", m).Msg("stream send message error")
	}
	if w.measurement == nil {
		return nil
	}

	if !w.measurement.IsFirstDocSent() {
		// The request is not tagged yet with db and collection, need to do it on the first message
		project, branch, coll := request.GetProjectAndBranchAndColl(m)
		reqMetadata, err := request.GetRequestMetadataFromContext(w.WrappedContext)
		if err != nil {
			if config.DefaultConfig.Metrics.DebugMessages {
				log.Debug().Str("error", err.Error()).Msg("error while getting request metadata, not measuring")
			}
			return nil
		}
		reqMetadata.SetProject(project)
		reqMetadata.SetBranch(branch)
		reqMetadata.SetCollection(coll)
		w.measurement.AddTags(map[string]string{
			"human": strconv.FormatBool(reqMetadata.IsHuman),
		})
		// Sets the flag so first time processing does not happen with subsequent documents
		w.measurement.MarkFirstDocSent()
		w.measurement.RecordDuration(metrics.RequestsRespTimeToFirstDoc, w.measurement.GetRequestOkTags())
		reqMetadata.SaveToContext(w.WrappedContext)
		w.WrappedContext = reqMetadata.SaveToContext(w.WrappedContext)

		w.measurement.AddProjectBranchCollTags(project, branch, coll)
	}
	w.measurement.AddSentBytes(proto.Size(m.(proto.Message)))
	return nil
}