fix(remote): race during flush (#154)

Author: korniltsev

upstream/remote/remote.go

@@ -30,15 +30,16 @@ const (
 )
 
 type Remote struct {
+	mu     sync.Mutex
 	cfg    Config
-	jobs   chan *upstream.UploadJob
+	jobs   chan job
 	client HTTPClient
 	logger Logger
 
 	done chan struct{}
 	wg   sync.WaitGroup
 
-	flushWG sync.WaitGroup
+	flushWG *sync.WaitGroup
 }
 
 type HTTPClient interface {
@@ -69,7 +70,7 @@ type Logger interface {
 func NewRemote(cfg Config) (*Remote, error) {
 	r := &Remote{
 		cfg:  cfg,
-		jobs: make(chan *upstream.UploadJob, 20),
+		jobs: make(chan job, 20),
 		client: &http.Client{
 			Transport: &http.Transport{
 				MaxConnsPerHost: cfg.Threads,
@@ -84,8 +85,9 @@ func NewRemote(cfg Config) (*Remote, error) {
 			},
 			Timeout: cfg.Timeout,
 		},
-		logger: cfg.Logger,
-		done:   make(chan struct{}),
+		logger:  cfg.Logger,
+		done:    make(chan struct{}),
+		flushWG: new(sync.WaitGroup),
 	}
 	if cfg.HTTPClient != nil {
 		r.client = cfg.HTTPClient
@@ -121,21 +123,28 @@ func (r *Remote) Stop() {
 	r.wg.Wait()
 }
 
-func (r *Remote) Upload(j *upstream.UploadJob) {
+func (r *Remote) Upload(uj *upstream.UploadJob) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	r.flushWG.Add(1)
+	j := job{
+		upload: uj,
+		flush:  r.flushWG,
+	}
 	select {
 	case r.jobs <- j:
 	default:
-		r.flushWG.Done()
+		j.flush.Done()
 		r.logger.Errorf("remote upload queue is full, dropping a profile job")
 	}
 }
 
 func (r *Remote) Flush() {
-	if r.done == nil {
-		return
-	}
-	r.flushWG.Wait()
+	r.mu.Lock()
+	flush := r.flushWG
+	r.flushWG = new(sync.WaitGroup)
+	r.mu.Unlock()
+	flush.Wait()
 }
 
 func (r *Remote) uploadProfile(j *upstream.UploadJob) error {
@@ -174,7 +183,6 @@ func (r *Remote) uploadProfile(j *upstream.UploadJob) error {
 
 	q := u.Query()
 	q.Set("name", j.Name)
-	// TODO: I think these should be renamed to startTime / endTime
 	q.Set("from", strconv.FormatInt(j.StartTime.UnixNano(), 10))
 	q.Set("until", strconv.FormatInt(j.EndTime.UnixNano(), 10))
 	q.Set("spyName", j.SpyName)
@@ -238,9 +246,9 @@ func (r *Remote) handleJobs() {
 		case <-r.done:
 			r.wg.Done()
 			return
-		case job := <-r.jobs:
-			r.safeUpload(job)
-			r.flushWG.Done()
+		case j := <-r.jobs:
+			r.safeUpload(j.upload)
+			j.flush.Done()
 		}
 	}
 }
@@ -262,3 +270,8 @@ func (r *Remote) safeUpload(job *upstream.UploadJob) {
 		r.logger.Errorf("upload profile: %v", err)
 	}
 }
+
+type job struct {
+	upload *upstream.UploadJob
+	flush  *sync.WaitGroup
+}

upstream/remote/remote_test.go

@@ -2,8 +2,11 @@ package remote
 
 import (
 	"bytes"
+	"fmt"
+	"github.com/stretchr/testify/require"
 	"io"
 	"net/http"
+	"sync"
 	"testing"
 	"time"
 
@@ -98,5 +101,62 @@ type MockHTTPClient struct {
 
 func (m *MockHTTPClient) Do(req *http.Request) (*http.Response, error) {
 	args := m.Called(req)
-	return args.Get(0).(*http.Response), args.Error(1)
+	err := args.Error(1)
+	a0 := args.Get(0)
+	switch typed := a0.(type) {
+	case *http.Response:
+		return typed, err
+	case func() *http.Response:
+		return typed(), err
+	default:
+		return nil, fmt.Errorf("unknown mock arg type arg %+v %w", a0, err)
+	}
+}
+
+func TestConcurrentUploadFlushRace(t *testing.T) {
+	mockClient := new(MockHTTPClient)
+	mockClient.On("Do", mock.Anything).Return(func() *http.Response {
+		return &http.Response{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewBufferString("OK")),
+		}
+	}, nil)
+	r, err := NewRemote(Config{
+		Threads:    2,
+		Logger:     testutil.NewTestLogger(),
+		HTTPClient: mockClient,
+	})
+	require.NoError(t, err)
+	r.Start()
+	defer r.Stop()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	loop := func(f func()) {
+		timeout := time.After(10 * time.Millisecond)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-timeout:
+					return
+				default:
+					f()
+				}
+			}
+		}()
+	}
+	loop(func() {
+		r.Upload(newJob("job1"))
+	})
+	loop(func() {
+		r.Flush()
+	})
+	wg.Wait()
+}
+
+func newJob(name string) *upstream.UploadJob {
+	return &upstream.UploadJob{
+		Name: name,
+	}
 }