main.go

package main

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"

	"cloud.google.com/go/storage"
	cloudtrace "github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace"
	"github.com/googleapis/gax-go/v2"
	"go.opentelemetry.io/contrib/detectors/gcp"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.26.0"
	"go.opentelemetry.io/otel/trace"
	"golang.org/x/oauth2"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
	"google.golang.org/grpc"

	"github.com/google/uuid"
)

/*

Sample command to run from a GCE VM:

go run . --client-protocol grpc --bucket thrivikramks_fio_test_us_west4 --read-folder-prefix "experiment.0/"

*/

var (
	maxConnsPerHost     = 0
	maxIdleConnsPerHost = 100
	serviceName         = "gcs-client-test"

	mountID = getRunMountId()

	// MiB means 1024 KiB.
	MiB = int64(1024 * 1024)

	grpcConnPoolSize = flag.Int("grpc-conn-pool-size", 1, "grpc connection pool size")

	maxRetryDuration = 30 * time.Second

	retryMultiplier = 2.0

	bucketName = flag.String("bucket", "kislayk_europe_west4", "GCS bucket name.")

	clientProtocol = flag.String("client-protocol", "http", "Network protocol.")

	// Object name = objectNamePrefix + {thread_id} + objectNameSuffix
	objectNamePrefix       = flag.String("obj-prefix", "1GB/experiment.", "Object prefix")
	objectNameSuffix       = flag.String("obj-suffix", ".0", "Object suffix")
	readFolderNamePrefix   = flag.String("read-folder-prefix", "", "Read folder prefix")
	disableOtelGRPCHandler = flag.Bool("disable-otel-grpc-handler", false, "Disable otel grpc handler")
)

// CreateHTTPClient create http storage client.
func CreateHTTPClient(ctx context.Context) (client *storage.Client, err error) {
	var transport *http.Transport
	// Using http1 makes the client more performant.
	transport = &http.Transport{
		MaxConnsPerHost:     maxConnsPerHost,
		MaxIdleConnsPerHost: maxIdleConnsPerHost,
		// This disables HTTP/2 in transport.
		TLSNextProto: make(
			map[string]func(string, *tls.Conn) http.RoundTripper,
		),
	}

	tokenSource, err := GetTokenSource(ctx, "")
	if err != nil {
		return nil, fmt.Errorf("while generating tokenSource, %v", err)
	}

	// Custom http client for Go Client.
	httpClient := &http.Client{
		Transport: &userAgentRoundTripper{
			wrapped: &oauth2.Transport{
				Base:   transport,
				Source: tokenSource,
			},
			UserAgent: "prince",
		},
		Timeout: 0,
	}
	return storage.NewClient(ctx, option.WithHTTPClient(httpClient))
}

// CreateGrpcClient creates grpc client.
func CreateGrpcClient(ctx context.Context) (client *storage.Client, err error) {
	tokenSource, err := GetTokenSource(ctx, "")
	if err != nil {
		return nil, err
	}

	var clientOptions []option.ClientOption
	clientOptions = append(clientOptions,
		option.WithGRPCConnectionPool(*grpcConnPoolSize),
		option.WithTokenSource(tokenSource),
		storage.WithDisabledClientMetrics(),
	)

	if !*disableOtelGRPCHandler {
		clientOptions = append(clientOptions, option.WithGRPCDialOption(grpc.WithStatsHandler(otelgrpc.NewClientHandler())))
	}

	return storage.NewGRPCClient(ctx, clientOptions...)
}

// ReadObject creates reader object corresponding to workerID with the help of bucketHandle.
func ReadObject(ctx context.Context, workerID int, bucketHandle *storage.BucketHandle) (bytesRead int64, err error) {
	objectName := *objectNamePrefix + strconv.Itoa(workerID) + *objectNameSuffix

	select {
	case <-ctx.Done():
		return 0, nil
	default:
		object := bucketHandle.Object(objectName)
		rc, err := object.NewReader(ctx)
		if err != nil {
			return 0, fmt.Errorf("while creating reader object: %v", err)
		}
		defer rc.Close()

		// Calls Reader.WriteTo implicitly.
		count, err := io.Copy(io.Discard, rc)
		bytesRead += count
		if err != nil {
			return bytesRead, fmt.Errorf("while reading and discarding content: %v", err)
		}
		return bytesRead, nil
	}
}

func ListObjects(ctx context.Context, bucketHandle *storage.BucketHandle) (err error) {
	itr := bucketHandle.Objects(ctx, &storage.Query{
		Prefix: *readFolderNamePrefix,
	})
	for {
		objAttrs, err := itr.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return fmt.Errorf("while listing objects: %v", err)
		}
		fmt.Println(objAttrs.Name)
	}

	return nil
}

func getResource(ctx context.Context, mountID string) (*resource.Resource, error) {
	fmt.Printf("Logging the mount ID: %v\n", mountID)
	return resource.New(ctx,
		resource.WithDetectors(gcp.NewDetector()),
		resource.WithAttributes(
			semconv.ServiceName(serviceName),
			semconv.ServiceInstanceID(mountID),
		),
	)
}

func getRunMountId() string {
	uuid := uuid.New()
	return uuid.String()
}

func newGCPCloudTraceExporter(ctx context.Context, mountID string) (*sdktrace.TracerProvider, ShutdownFn, error) {
	var traceOptions []cloudtrace.Option

	exporter, err := cloudtrace.New(traceOptions...)

	if err != nil {
		return nil, nil, err
	}
	res, err := getResource(ctx, mountID)
	if err != nil {
		return nil, nil, err
	}

	tp := sdktrace.NewTracerProvider(
		sdktrace.WithSpanProcessor(sdktrace.NewBatchSpanProcessor(exporter)),
		sdktrace.WithResource(res), sdktrace.WithSampler(sdktrace.AlwaysSample()))

	return tp, tp.Shutdown, nil
}

func newTraceProvider(ctx context.Context, mountID string) (trace.TracerProvider, ShutdownFn, error) {
	return newGCPCloudTraceExporter(ctx, mountID)
}

type ShutdownFn func(ctx context.Context) error

// SetupTracing bootstraps the OpenTelemetry tracing pipeline.
func SetupTracing(ctx context.Context) ShutdownFn {
	tp, shutdown, err := newTraceProvider(ctx, mountID)

	otel.SetTextMapPropagator(propagation.TraceContext{})

	if err != nil {
		fmt.Printf("error occurred while setting up tracing: %v", err)
		return nil
	}
	otel.SetTracerProvider(tp)
	return shutdown
}

func ReadDirTracer(ctx context.Context, bucketHandle *storage.BucketHandle) (err error) {
	tracer := otel.Tracer("test_gcsfuse")
	ctx, span := tracer.Start(ctx, "ReadDirMimic")
	defer span.End()
	err = ListObjects(ctx, bucketHandle)
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
	}
	return
}

func main() {
	flag.Parse()
	fmt.Printf("Workload start time: %s\n", time.Now().String())
	ctx := context.Background()

	var client *storage.Client
	var err error
	protocol := ""
	if *clientProtocol == "http" {
		protocol = "http"
		client, err = CreateHTTPClient(ctx)
	} else {
		protocol = "grpc"
		client, err = CreateGrpcClient(ctx)
	}
	fmt.Printf("Client Protocol is %v\n", protocol)

	if err != nil {
		fmt.Printf("while creating the client: %v", err)
		os.Exit(1)
	}

	client.SetRetry(
		storage.WithBackoff(gax.Backoff{
			Max:        maxRetryDuration,
			Multiplier: retryMultiplier,
		}),
		storage.WithPolicy(storage.RetryAlways))

	// assumes bucket already exist
	bucketHandle := client.Bucket(*bucketName)

	// Run the trace ListObjects workload
	newCtx := context.Background()
	traceShutdownFn := SetupTracing(newCtx)
	defer func() {
		traceShutdownFn(newCtx)
	}()
	ReadDirTracer(newCtx, bucketHandle)
}