Inspect GFE Latency in Spanner Go Client

main.go

// Copyright 2019 Google LLC
//
// 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
//
//     https://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.

// [START spanner_quickstart]

// Sample spanner_quickstart is a basic program that uses Cloud Spanner.
package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	spanner "cloud.google.com/go/spanner/apiv1"
	gax "github.com/googleapis/gax-go/v2"
	sppb "google.golang.org/genproto/googleapis/spanner/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	"contrib.go.opencensus.io/exporter/prometheus"
	"go.opencensus.io/metric"
	"go.opencensus.io/metric/metricdata"
	"go.opencensus.io/metric/metricproducer"
)

func main() {
	ctx := context.Background()

	// This database must exist.
	databaseName := "projects/testing/instances/test-instance/databases/test-db"

	client, err := spanner.NewClient(ctx)
	if err != nil {
		log.Fatalf("Failed to create client %v", err)
	}
	defer client.Close()

	gauge := setupOpenCensus()

	req := &sppb.CreateSessionRequest{Database: databaseName}
	session, err := client.CreateSession(ctx, req)
	if err != nil {
		log.Fatalf("CreateSession failed with %v", err)
	}

	req2 := &sppb.ExecuteSqlRequest{
		Session:   session.Name,
		Sql:       "SELECT * FROM Test ORDER BY ID ASC",
		QueryMode: sppb.ExecuteSqlRequest_PROFILE,
	}

	for i := 0; i < 20; i++ {
		start := time.Now()
		var md metadata.MD
		resultSet, err := client.ExecuteSql(ctx, req2, gax.WithGRPCOptions(grpc.Header(&md)))
		if err != nil {
			log.Fatalf("ExecuteSql failed with %v", err)
		}
		fmt.Println("Row", resultSet.GetRows())
		recordMetrics(md, resultSet, start, gauge)
		time.Sleep(10 * time.Second)
	}
}

func setupOpenCensus() *metric.Int64Gauge {
	// Create metric registry and register it with global producer manager.
	r := metric.NewRegistry()
	metricproducer.GlobalManager().AddProducer(r)

	pe, err := prometheus.NewExporter(prometheus.Options{
		Namespace: "demo",
	})
	if err != nil {
		log.Fatalf("Failed to create the Prometheus exporter: %v", err)
	}
	go func() {
		mux := http.NewServeMux()
		mux.Handle("/metrics", pe)
		if err := http.ListenAndServe(":8888", mux); err != nil {
			log.Fatalf("Failed to run Prometheus /metrics endpoint: %v", err)
		}
	}()

	return createGauge(r, "elapsed_time", "Elapsed time")
}

func createGauge(r *metric.Registry, name, description string) *metric.Int64Gauge {
	gauge, err := r.AddInt64Gauge(
		name,
		metric.WithDescription(description),
		metric.WithUnit(metricdata.UnitMilliseconds),
		metric.WithLabelKeys("stage"),
		metric.WithConstLabel(map[metricdata.LabelKey]metricdata.LabelValue{{Key: "method"}: metricdata.NewLabelValue("ExecuteSql")}),
	)
	if err != nil {
		log.Fatalf("error creating %s, error %v\n", name, err)
	}
	return gauge
}

func recordMetrics(md metadata.MD, resultSet *sppb.ResultSet, start time.Time, gauge *metric.Int64Gauge) {
	clientElapsed := time.Since(start).Milliseconds()
	set("client", int64(clientElapsed), gauge)

	srvTiming := md.Get("server-timing")[0]
	gfeLtcy, _ := strconv.Atoi(strings.TrimPrefix(srvTiming, "gfet4t7; dur="))
	set("gfe", int64(gfeLtcy), gauge)

	stats := resultSet.GetStats()
	fmt.Println("query elaspsed time", stats.QueryStats.Fields["elapsed_time"])
	spanFELtcy := stats.QueryStats.Fields["elapsed_time"].GetStringValue()
	spanFELtcyMS, _ := strconv.ParseFloat(strings.TrimSuffix(spanFELtcy, " msecs"), 32)
	fmt.Println("query elaspsed time ms", spanFELtcyMS)
	set("spanner_query", int64(spanFELtcyMS), gauge)
}

func set(labelVal string, latency int64, gauge *metric.Int64Gauge) {
	if entry, err := gauge.GetEntry(metricdata.NewLabelValue(labelVal)); err != nil {
		log.Fatalf("Error %v when GetEntry", err)
	} else {
		entry.Set(int64(latency))
	}
}

// [END spanner_quickstart]