From 25a887aa8d6be7fbbe60d3eda4a34ba82249ba53 Mon Sep 17 00:00:00 2001
From: Chris Cotter <cjcotter@google.com>
Date: Wed, 2 Aug 2023 21:55:56 +0000
Subject: [PATCH] feat(storage): add trace span to Writer

Adds a trace span that covers a Writer from open to close.
Nested spans already cover individual RPCs that the Writer
makes.

Fixes #6144
---
 storage/storage.go | 1 +
 storage/writer.go  | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/storage/storage.go b/storage/storage.go
index ba485ddcdacb..320439da5432 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -1036,6 +1036,7 @@ func (o *ObjectHandle) ReadCompressed(compressed bool) *ObjectHandle {
 // It is the caller's responsibility to call Close when writing is done. To
 // stop writing without saving the data, cancel the context.
 func (o *ObjectHandle) NewWriter(ctx context.Context) *Writer {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/storage.Object.Writer")
 	return &Writer{
 		ctx:         ctx,
 		o:           o,
diff --git a/storage/writer.go b/storage/writer.go
index f06c31162670..849c5a89ea85 100644
--- a/storage/writer.go
+++ b/storage/writer.go
@@ -15,6 +15,7 @@
 package storage
 
 import (
+	"cloud.google.com/go/internal/trace"
 	"context"
 	"errors"
 	"fmt"
@@ -163,6 +164,7 @@ func (w *Writer) Close() error {
 	<-w.donec
 	w.mu.Lock()
 	defer w.mu.Unlock()
+	trace.EndSpan(w.ctx, w.err)
 	return w.err
 }