[DPL-1341] Fix concurrency bug

.github/workflows/go.yml

@@ -11,12 +11,12 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
 
     - name: Set up Go
-      uses: actions/setup-go@v2
+      uses: actions/setup-go@v5
       with:
-        go-version: 1.17
+        go-version: 1.22.5
 
     - name: Build
       run: go build

.gitignore

@@ -0,0 +1 @@
+.idea

go.mod

@@ -1,10 +1,14 @@
 module regproxy2
 
-go 1.17
+go 1.22.5
+
+require (
+	go.mercari.io/go-dnscache v0.0.0-20210517095825-88b046eb94f2
+	go.uber.org/zap v1.19.1
+	golang.org/x/sync v0.7.0
+)
 
 require (
-	go.mercari.io/go-dnscache v0.0.0-20210517095825-88b046eb94f2 // indirect
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.7.0 // indirect
-	go.uber.org/zap v1.19.1 // indirect
 )

go.sum

@@ -31,6 +31,8 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL
 golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M=
+golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

main.go

@@ -11,6 +11,7 @@ import (
 	"net/url"
 	"os"
 	"strconv"
+	"sync"
 	"time"
 
 	dnscache "go.mercari.io/go-dnscache"
@@ -34,6 +35,7 @@ func errResp(resp http.ResponseWriter, e error) {
 type RegProxy struct {
 	client    *http.Client
 	upstreams map[string]*url.URL
+	writeLock sync.Mutex
 	handler   http.Handler
 }
 
@@ -130,6 +132,9 @@ func (self *RegProxy) register(resp http.ResponseWriter, req *http.Request) {
 		badRequest(resp, err.Error())
 		return
 	}
+
+	self.writeLock.Lock()
+	defer self.writeLock.Unlock()
 	log.Printf("Adding upstream %v", q)
 	self.upstreams[q.Name] = upstream
 	resp.WriteHeader(204)

main_test.go

@@ -2,7 +2,10 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"fmt"
+	"golang.org/x/sync/errgroup"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -63,6 +66,32 @@ func TestRegister(t *testing.T) {
 	}
 }
 
+func TestConcurrentRegister(t *testing.T) {
+	errorGroup, _ := errgroup.WithContext(context.Background())
+
+	withRegProxy(t, func(url string, t *testing.T) {
+		for i := 0; i < 100; i++ {
+			request := fmt.Sprintf("{\"name\":\"foo-%d\",\"callback\":\"baz\"}", i)
+			errorGroup.Go(func() error {
+				r, err := http.Post(url+"/register", "application/json", bytes.NewReader([]byte(request)))
+				if err != nil {
+					return err
+				}
+				if r.StatusCode != 204 {
+					return fmt.Errorf("Wrong status code from /register %d expected 204", r.StatusCode)
+				}
+
+				return nil
+			})
+		}
+
+		resultingError := errorGroup.Wait()
+		if resultingError != nil {
+			t.Fatal(resultingError)
+		}
+	})
+}
+
 func register(url string, u upstream, t *testing.T) {
 	b, _ := json.Marshal(u)
 	r, err := http.Post(url+"/register", "application/json", bytes.NewReader(b))