profiler: fix start to allow retry until success
This change allows profiler agent Start to be called multiple times
until it succeeds. After a successful start, it will not be called
again.
@@ -47,6 +47,7 @@ import (
"runtime"
"runtime/pprof"
"sync"
+ "sync/atomic"
"time"
gcemd "cloud.google.com/go/compute/metadata"
@@ -67,7 +68,7 @@ import (
var (
config Config
- startOnce sync.Once
+ startOnce allowUntilSuccess
mutexEnabled bool
// The functions below are stubbed to be overrideable for testing.
getProjectID = gcemd.ProjectID
@@ -173,13 +174,40 @@ type Config struct {
// initializating and starting of the agent.
var startError error
+// allowUntilSuccess is an object that will perform action till
+// it succeeds once.
+// This is a modified form of Go's sync.Once
+type allowUntilSuccess struct {
+ m sync.Mutex
+ done uint32
+}
+
+// do calls function f only if it hasnt returned nil previously.
+// Once f returns nil, do will not call function f any more.
+// This is a modified form of Go's sync.Once.Do
+func (o *allowUntilSuccess) do(f func() error) (err error) {
+ if atomic.LoadUint32(&o.done) == 1 {
+ log.Printf("profiler.Start() called again after it was previously called")
+ return nil
+ }
+ // Slow-path.
+ o.m.Lock()
+ defer o.m.Unlock()
+ if o.done == 0 {
+ if err = f(); err == nil {
+ atomic.StoreUint32(&o.done, 1)
+ }
+ }
+ return err
+}
+
// Start starts a goroutine to collect and upload profiles. The
// caller must provide the service string in the config. See
// Config for details. Start should only be called once. Any
// additional calls will be ignored.
func Start(cfg Config, options ...option.ClientOption) error {
- startOnce.Do(func() {
- startError = start(cfg, options...)
+ startError = startOnce.do(func() error {
+ return start(cfg, options...)
})
return startError
}
