restful_middleware.go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 metrics

import (
	"context"
	"strconv"
)

func WrapRestfulContext(ctx context.Context, inputLength int64) context.Context {
	return context.WithValue(ctx, milvusStatsKey{}, &RPCStats{
		inboundPayloadSize: int(inputLength),
	})
}

func RecordRestfulMetrics(ctx context.Context, outputLength int64, observeOutbound bool) {
	if mstats := GetStats(ctx); mstats != nil {
		// all info set
		// set metrics with inbound size and related meta
		nodeIDValue := strconv.FormatInt(mstats.nodeID, 10)
		if mstats.inboundPayloadSize > 0 {
			ProxyReceiveBytes.WithLabelValues(
				nodeIDValue,
				mstats.inboundLabel, mstats.collectionName).Add(float64(mstats.inboundPayloadSize))
		}
		// set outbound payload size metrics
		if outputLength > 0 && observeOutbound {
			ProxyReadReqSendBytes.WithLabelValues(nodeIDValue).Add(float64(outputLength))
		}
	}
}