add tracing (#152)

Co-authored-by: James Kwon <[email protected]>
Comfy-Org · Jan 23, 2025 · 2a4b225 · 2a4b225
1 parent 6828234
commit 2a4b225
Show file tree

Hide file tree

Showing 16 changed files with 211 additions and 190 deletions.
diff --git a/gateways/algolia/algolia.go b/gateways/algolia/algolia.go
@@ -7,6 +7,7 @@ import (
 	"registry-backend/ent"
 	"registry-backend/entity"
 	"registry-backend/mapper"
+	"registry-backend/tracing"
 
 	"github.com/algolia/algoliasearch-client-go/v3/algolia/search"
 	"github.com/rs/zerolog/log"
@@ -49,6 +50,8 @@ func NewAlgoliaService(cfg *config.Config) (AlgoliaService, error) {
 
 // IndexNodes indexes the provided nodes in Algolia.
 func (a *algolia) IndexNodes(ctx context.Context, nodes ...*ent.Node) error {
+	defer tracing.TraceDefaultSegment(ctx, "AlgoliaService.IndexNodes")()
+
 	index := a.client.InitIndex("nodes_index")
 	objects := make([]entity.AlgoliaNode, len(nodes))
 
@@ -66,6 +69,8 @@ func (a *algolia) IndexNodes(ctx context.Context, nodes ...*ent.Node) error {
 
 // SearchNodes searches for nodes in Algolia matching the query.
 func (a *algolia) SearchNodes(ctx context.Context, query string, opts ...interface{}) (nodes []*ent.Node, err error) {
+	defer tracing.TraceDefaultSegment(ctx, "AlgoliaService.SearchNodes")()
+
 	index := a.client.InitIndex("nodes_index")
 	res, err := index.Search(query, opts...)
 	if err != nil {
@@ -84,6 +89,8 @@ func (a *algolia) SearchNodes(ctx context.Context, query string, opts ...interfa
 
 // DeleteNode deletes the specified node from Algolia.
 func (a *algolia) DeleteNode(ctx context.Context, node *ent.Node) error {
+	defer tracing.TraceDefaultSegment(ctx, "AlgoliaService.DeleteNode")()
+
 	index := a.client.InitIndex("nodes_index")
 	res, err := index.DeleteObject(node.ID)
 	if err != nil {
@@ -94,6 +101,8 @@ func (a *algolia) DeleteNode(ctx context.Context, node *ent.Node) error {
 
 // IndexNodeVersions implements AlgoliaService.
 func (a *algolia) IndexNodeVersions(ctx context.Context, nodes ...*ent.NodeVersion) error {
+	defer tracing.TraceDefaultSegment(ctx, "AlgoliaService.IndexNodeVersions")()
+
 	index := a.client.InitIndex("node_versions_index")
 	objects := make([]struct {
 		ObjectID string `json:"objectID"`
@@ -122,6 +131,8 @@ func (a *algolia) IndexNodeVersions(ctx context.Context, nodes ...*ent.NodeVersi
 
 // DeleteNodeVersions implements AlgoliaService.
 func (a *algolia) DeleteNodeVersions(ctx context.Context, nodes ...*ent.NodeVersion) error {
+	defer tracing.TraceDefaultSegment(ctx, "AlgoliaService.DeleteNodeVersions")()
+
 	index := a.client.InitIndex("node_versions_index")
 	ids := []string{}
 	for _, node := range nodes {
@@ -136,6 +147,8 @@ func (a *algolia) DeleteNodeVersions(ctx context.Context, nodes ...*ent.NodeVers
 
 // SearchNodeVersions implements AlgoliaService.
 func (a *algolia) SearchNodeVersions(ctx context.Context, query string, opts ...interface{}) ([]*ent.NodeVersion, error) {
+	defer tracing.TraceDefaultSegment(ctx, "AlgoliaService.SearchNodeVersions")()
+
 	index := a.client.InitIndex("node_versions_index")
 	res, err := index.Search(query, opts...)
 	if err != nil {

diff --git a/gateways/pubsub/pubsub.go b/gateways/pubsub/pubsub.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/url"
 	"registry-backend/config"
+	"registry-backend/tracing"
 	"strconv"
 	"strings"
 	"time"
@@ -47,6 +48,8 @@ func NewPubSubService(c *config.Config) (PubSubService, error) {
 
 // PublishNodePack implements PubSubService.
 func (p *pubsubimpl) PublishNodePack(ctx context.Context, storageURL string) (err error) {
+	defer tracing.TraceDefaultSegment(ctx, "PubSubService.PublishNodePack")()
+
 	u, err := url.Parse(storageURL)
 	if err != nil {
 		return fmt.Errorf("invalid storage URL: %w", err)

diff --git a/gateways/storage/files.go b/gateways/storage/files.go
@@ -8,9 +8,11 @@ import (
 	"os"
 	"time"
 
+	"registry-backend/config"
+	"registry-backend/tracing"
+
 	"cloud.google.com/go/storage"
 	"github.com/rs/zerolog/log"
-	"registry-backend/config"
 )
 
 // StorageService defines the interface for interacting with cloud storage.
@@ -52,6 +54,8 @@ func NewStorageService(cfg *config.Config) (StorageService, error) {
 
 // UploadFile uploads an object to GCP storage.
 func (s *storageService) UploadFile(ctx context.Context, bucket, object, filePath string) (string, error) {
+	defer tracing.TraceDefaultSegment(ctx, "StorageService.UploadFile")()
+
 	log.Ctx(ctx).Info().Msgf("Uploading %v to %v/%v.\n", filePath, bucket, object)
 
 	// Open local file
@@ -120,6 +124,8 @@ func (s *storageService) StreamFileUpload(w io.Writer, objectName, blob string)
 
 // GetFileUrl gets the public URL of a file from GCP storage.
 func (s *storageService) GetFileUrl(ctx context.Context, bucketName, objectPath string) (string, error) {
+	defer tracing.TraceDefaultSegment(ctx, "StorageService.GetFileUrl")()
+
 	// Get the public URL of a file in a bucket
 	attrs, err := s.client.Bucket(bucketName).Object(objectPath).Attrs(ctx)
 	if err != nil {

diff --git a/server/implementation/cicd.go b/server/implementation/cicd.go
@@ -8,6 +8,7 @@ import (
 	"registry-backend/ent/gitcommit"
 	"registry-backend/ent/schema"
 	"registry-backend/mapper"
+	"registry-backend/tracing"
 	"sort"
 	"strings"
 
@@ -17,6 +18,8 @@ import (
 )
 
 func (impl *DripStrictServerImplementation) GetGitcommit(ctx context.Context, request drip.GetGitcommitRequestObject) (drip.GetGitcommitResponseObject, error) {
+	defer tracing.TraceDefaultSegment(ctx, "DripStrictServerImplementation.GetGitcommit")()
+
 	var commitId uuid.UUID = uuid.Nil
 	if request.Params.CommitId != nil {
 		log.Ctx(ctx).Info().Msgf("getting commit data for %s", *request.Params.CommitId)
@@ -124,6 +127,7 @@ func (impl *DripStrictServerImplementation) GetGitcommit(ctx context.Context, re
 }
 
 func (impl *DripStrictServerImplementation) GetGitcommitsummary(ctx context.Context, request drip.GetGitcommitsummaryRequestObject) (drip.GetGitcommitsummaryResponseObject, error) {
+	defer tracing.TraceDefaultSegment(ctx, "DripStrictServerImplementation.GetGitcommitsummary")()
 	log.Ctx(ctx).Info().Msg("Getting git commit summary")
 
 	// Prep relevant vars
@@ -254,6 +258,8 @@ func (impl *DripStrictServerImplementation) GetGitcommitsummary(ctx context.Cont
 }
 
 func (impl *DripStrictServerImplementation) GetWorkflowResult(ctx context.Context, request drip.GetWorkflowResultRequestObject) (drip.GetWorkflowResultResponseObject, error) {
+	defer tracing.TraceDefaultSegment(ctx, "DripStrictServerImplementation.GetWorkflowResult")()
+
 	log.Ctx(ctx).Info().Msgf("Getting workflow result with ID %s", request.WorkflowResultId)
 	workflowId := uuid.MustParse(request.WorkflowResultId)
 	workflow, err := impl.Client.CIWorkflowResult.Query().WithGitcommit().WithStorageFile().Where(ciworkflowresult.IDEQ(workflowId)).First(ctx)
@@ -278,6 +284,8 @@ func (impl *DripStrictServerImplementation) GetWorkflowResult(ctx context.Contex
 }
 
 func (impl *DripStrictServerImplementation) GetBranch(ctx context.Context, request drip.GetBranchRequestObject) (drip.GetBranchResponseObject, error) {
+	defer tracing.TraceDefaultSegment(ctx, "DripStrictServerImplementation.GetBranch")()
+
 	repoNameFilter := strings.ToLower(request.Params.RepoName)
 
 	branches, err := impl.Client.GitCommit.
@@ -295,6 +303,8 @@ func (impl *DripStrictServerImplementation) GetBranch(ctx context.Context, reque
 }
 
 func (impl *DripStrictServerImplementation) PostUploadArtifact(ctx context.Context, request drip.PostUploadArtifactRequestObject) (drip.PostUploadArtifactResponseObject, error) {
+	defer tracing.TraceDefaultSegment(ctx, "DripStrictServerImplementation.PostUploadArtifact")()
+
 	err := impl.ComfyCIService.ProcessCIRequest(ctx, impl.Client, &request)
 	if err != nil {
 		log.Ctx(ctx).Error().Msgf("Error processing CI request w/ err: %v", err)