Add opentelemetry instrumentation

go.mod

@@ -24,14 +24,20 @@ require (
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.11.0
 	github.com/stretchr/testify v1.8.4
+	go.opentelemetry.io/otel v1.11.1
+	go.opentelemetry.io/otel/exporters/jaeger v1.11.1
+	go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.11.1
+	go.opentelemetry.io/otel/sdk v1.11.1
+	go.opentelemetry.io/otel/trace v1.11.1
 	golang.org/x/time v0.1.0
 	golang.org/x/tools v0.6.0
 	google.golang.org/grpc v1.56.1
 	google.golang.org/protobuf v1.30.0
 	gorm.io/gorm v1.22.4
-	k8s.io/api v0.20.2
-	k8s.io/apimachinery v0.20.2
-	k8s.io/client-go v0.0.0-20210217172142-7279fc64d847
+	k8s.io/api v0.24.0
+	k8s.io/apimachinery v0.24.0
+	k8s.io/client-go v0.24.0
+	sigs.k8s.io/controller-runtime v0.12.1
 )
 
 require (
@@ -50,15 +56,24 @@ require (
 	github.com/Azure/go-autorest/autorest/date v0.3.0 // indirect
 	github.com/Azure/go-autorest/logger v0.2.1 // indirect
 	github.com/Azure/go-autorest/tracing v0.6.0 // indirect
+	github.com/PuerkitoBio/purell v1.1.1 // indirect
+	github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash v1.1.0 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/go-logr/logr v0.4.0 // indirect
+	github.com/emicklei/go-restful v2.9.5+incompatible // indirect
+	github.com/evanphx/json-patch v4.12.0+incompatible // indirect
+	github.com/go-logr/logr v1.2.3 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
+	github.com/go-openapi/jsonpointer v0.19.5 // indirect
+	github.com/go-openapi/jsonreference v0.19.5 // indirect
+	github.com/go-openapi/swag v0.19.14 // indirect
 	github.com/gofrs/uuid v4.2.0+incompatible // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang-jwt/jwt/v4 v4.4.1 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/google/gnostic v0.5.7-v3refs // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/google/gofuzz v1.1.0 // indirect
 	github.com/google/uuid v1.3.0 // indirect
@@ -67,13 +82,15 @@ require (
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
-	github.com/kr/text v0.2.0 // indirect
+	github.com/mailru/easyjson v0.7.6 // indirect
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
-	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/ncw/swift v1.0.53 // indirect
 	github.com/pelletier/go-toml v1.9.4 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.0-beta.8 // indirect
@@ -92,6 +109,7 @@ require (
 	golang.org/x/net v0.9.0 // indirect
 	golang.org/x/oauth2 v0.7.0 // indirect
 	golang.org/x/sys v0.7.0 // indirect
+	golang.org/x/term v0.7.0 // indirect
 	golang.org/x/text v0.9.0 // indirect
 	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 	google.golang.org/api v0.114.0 // indirect
@@ -101,6 +119,10 @@ require (
 	gopkg.in/ini.v1 v1.66.4 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
-	k8s.io/klog/v2 v2.5.0 // indirect
-	sigs.k8s.io/structured-merge-diff/v4 v4.0.3 // indirect
+	k8s.io/klog/v2 v2.60.1 // indirect
+	k8s.io/kube-openapi v0.0.0-20220328201542-3ee0da9b0b42 // indirect
+	k8s.io/utils v0.0.0-20220210201930-3a6ce19ff2f9 // indirect
+	sigs.k8s.io/json v0.0.0-20211208200746-9f7c6b3444d2 // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.2.1 // indirect
+	sigs.k8s.io/yaml v1.3.0 // indirect
 )

storage/cached_rawstore.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/flyteorg/flytestdlib/errors"
+	"github.com/flyteorg/flytestdlib/telemetryutils"
 
 	"github.com/coocood/freecache"
 	"github.com/prometheus/client_golang/prometheus"
@@ -49,6 +50,9 @@ func (s *cachedRawStore) Head(ctx context.Context, reference DataReference) (Met
 
 // ReadRaw retrieves a byte array from the Blob store or an error
 func (s *cachedRawStore) ReadRaw(ctx context.Context, reference DataReference) (io.ReadCloser, error) {
+	ctx, span := telemetryutils.NewSpan(ctx, telemetryutils.BlobstoreClientTracer, "flytestdlib.storage.cachedRawStore/ReadRaw")
+	defer span.End()
+
 	key := []byte(reference)
 	if oRaw, err := s.cache.Get(key); err == nil {
 		// Found, Cache hit
@@ -84,6 +88,9 @@ func (s *cachedRawStore) ReadRaw(ctx context.Context, reference DataReference) (
 
 // WriteRaw stores a raw byte array.
 func (s *cachedRawStore) WriteRaw(ctx context.Context, reference DataReference, size int64, opts Options, raw io.Reader) error {
+	ctx, span := telemetryutils.NewSpan(ctx, telemetryutils.BlobstoreClientTracer, "flytestdlib.storage.cachedRawStore/WriteRaw")
+	defer span.End()
+
 	var buf bytes.Buffer
 	teeReader := io.TeeReader(raw, &buf)
 	err := s.RawStore.WriteRaw(ctx, reference, size, opts, teeReader)

storage/protobuf_store.go

@@ -13,6 +13,7 @@ import (
 	"github.com/flyteorg/flytestdlib/ioutils"
 	"github.com/flyteorg/flytestdlib/logger"
 	"github.com/flyteorg/flytestdlib/promutils"
+	"github.com/flyteorg/flytestdlib/telemetryutils"
 )
 
 type protoMetrics struct {
@@ -32,6 +33,9 @@ type DefaultProtobufStore struct {
 }
 
 func (s DefaultProtobufStore) ReadProtobuf(ctx context.Context, reference DataReference, msg proto.Message) error {
+	ctx, span := telemetryutils.NewSpan(ctx, telemetryutils.BlobstoreClientTracer, "flytestdlib.storage.DefaultProtobufStore/ReadProtobuf")
+	defer span.End()
+
 	rc, err := s.ReadRaw(ctx, reference)
 	if err != nil && !IsFailedWriteToCache(err) {
 		logger.Errorf(ctx, "Failed to read from the raw store [%s] Error: %v", reference, err)
@@ -63,6 +67,9 @@ func (s DefaultProtobufStore) ReadProtobuf(ctx context.Context, reference DataRe
 }
 
 func (s DefaultProtobufStore) WriteProtobuf(ctx context.Context, reference DataReference, opts Options, msg proto.Message) error {
+	ctx, span := telemetryutils.NewSpan(ctx, telemetryutils.BlobstoreClientTracer, "flytestdlib.storage.DefaultProtobufStore/WriteProtobuf")
+	defer span.End()
+
 	t := s.metrics.MarshalTime.Start()
 	raw, err := proto.Marshal(msg)
 	t.Stop()

telemetryutils/config.go

@@ -0,0 +1,52 @@
+package telemetryutils
+
+import (
+	"context"
+
+	"github.com/flyteorg/flytestdlib/config"
+	"github.com/flyteorg/flytestdlib/logger"
+)
+
+//go:generate pflags Config --default-var=defaultConfig
+
+type Type = string
+
+const configSectionKey = "telemetry"
+
+var (
+	ConfigSection = config.MustRegisterSection(configSectionKey, defaultConfig)
+	defaultConfig = &Config{
+		FileConfig: FileConfig{
+			Enabled:  false,
+			Filename: "/tmp/trace.txt",
+		},
+		JaegerConfig: JaegerConfig{
+			Enabled:  false,
+			Endpoint: "http://localhost:14268/api/traces",
+		},
+	}
+)
+
+type Config struct {
+	FileConfig   FileConfig   `json:"file" pflag:",Configuration for exporting telemetry traces to a file"`
+	JaegerConfig JaegerConfig `json:"jaeger" pflag:",Configuration for exporting telemetry traces to a jaeger"`
+}
+
+type FileConfig struct {
+	Enabled  bool   `json:"enabled" pflag:",Set to true to enable the file exporter"`
+	Filename string `json:"filename" pflag:",Filename to store exported telemetry traces"`
+}
+
+type JaegerConfig struct {
+	Enabled  bool   `json:"enabled" pflag:",Set to true to enable the jaeger exporter"`
+	Endpoint string `json:"endpoint" pflag:",Endpoint for the jaeger telemtry trace ingestor"`
+}
+
+func GetConfig() *Config {
+	if c, ok := ConfigSection.GetConfig().(*Config); ok {
+		return c
+	}
+
+	logger.Warnf(context.TODO(), "Failed to retrieve config section [%v].", configSectionKey)
+	return nil
+}