Skip to content

Commit

Permalink
test: add RWLock to fix data race in MockDBAdapter
Browse files Browse the repository at this point in the history
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
  • Loading branch information
Harry-Chen committed Jan 11, 2025
1 parent 3ad551f commit 15e87a5
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions manager/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"math/rand"
"net/http"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
Expand Down Expand Up @@ -424,32 +425,41 @@ func TestHTTPServer(t *testing.T) {
type mockDBAdapter struct {
workerStore map[string]WorkerStatus
statusStore map[string]MirrorStatus
workerLock sync.RWMutex
statusLock sync.RWMutex
}

func (b *mockDBAdapter) Init() error {
return nil
}

func (b *mockDBAdapter) ListWorkers() ([]WorkerStatus, error) {
b.workerLock.RLock()
workers := make([]WorkerStatus, len(b.workerStore))
idx := 0
for _, w := range b.workerStore {
workers[idx] = w
idx++
}
b.workerLock.RUnlock()
return workers, nil
}

func (b *mockDBAdapter) GetWorker(workerID string) (WorkerStatus, error) {
b.workerLock.RLock()
defer b.workerLock.RUnlock()
w, ok := b.workerStore[workerID]

if !ok {
return WorkerStatus{}, fmt.Errorf("invalid workerId")
}
return w, nil
}

func (b *mockDBAdapter) DeleteWorker(workerID string) error {
b.workerLock.Lock()
delete(b.workerStore, workerID)
b.workerLock.Unlock()
return nil
}

Expand All @@ -458,7 +468,9 @@ func (b *mockDBAdapter) CreateWorker(w WorkerStatus) (WorkerStatus, error) {
// if ok {
// return workerStatus{}, fmt.Errorf("duplicate worker name")
// }
b.workerLock.Lock()
b.workerStore[w.ID] = w
b.workerLock.Unlock()
return w, nil
}

Expand All @@ -473,7 +485,9 @@ func (b *mockDBAdapter) RefreshWorker(workerID string) (w WorkerStatus, err erro

func (b *mockDBAdapter) GetMirrorStatus(workerID, mirrorID string) (MirrorStatus, error) {
id := mirrorID + "/" + workerID
b.statusLock.RLock()
status, ok := b.statusStore[id]
b.statusLock.RUnlock()
if !ok {
return MirrorStatus{}, fmt.Errorf("no mirror %s exists in worker %s", mirrorID, workerID)
}
Expand All @@ -487,7 +501,9 @@ func (b *mockDBAdapter) UpdateMirrorStatus(workerID, mirrorID string, status Mir
// }

id := mirrorID + "/" + workerID
b.statusLock.Lock()
b.statusStore[id] = status
b.statusLock.Unlock()
return status, nil
}

Expand All @@ -497,19 +513,23 @@ func (b *mockDBAdapter) ListMirrorStatus(workerID string) ([]MirrorStatus, error
if workerID == _magicBadWorkerID {
return []MirrorStatus{}, fmt.Errorf("database fail")
}
b.statusLock.RLock()
for k, v := range b.statusStore {
if wID := strings.Split(k, "/")[1]; wID == workerID {
mirrorStatusList = append(mirrorStatusList, v)
}
}
b.statusLock.RUnlock()
return mirrorStatusList, nil
}

func (b *mockDBAdapter) ListAllMirrorStatus() ([]MirrorStatus, error) {
var mirrorStatusList []MirrorStatus
b.statusLock.RLock()
for _, v := range b.statusStore {
mirrorStatusList = append(mirrorStatusList, v)
}
b.statusLock.RUnlock()
return mirrorStatusList, nil
}

Expand Down

0 comments on commit 15e87a5

Please sign in to comment.