From e5b1f8e54a00370018f53737563d8e0e61da42c9 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Fri, 2 Aug 2024 16:12:37 +0200
Subject: vm/dispatcher: fix data races

It was possible for poolInstance.reserve() and free() to race with
instance restart in Pool.Loop().

Add more locking to poolInstance. Remove locks in one case where it was
excessive.
---
 vm/dispatcher/pool.go | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

(limited to 'vm/dispatcher')

diff --git a/vm/dispatcher/pool.go b/vm/dispatcher/pool.go
index 7ad934cfb..d13812172 100644
--- a/vm/dispatcher/pool.go
+++ b/vm/dispatcher/pool.go
@@ -77,10 +77,7 @@ func (p *Pool[T]) runInstance(ctx context.Context, inst *poolInstance[T]) {
 
 	log.Logf(2, "pool: booting instance %d", inst.idx)
 
-	p.mu.Lock()
-	// Avoid races with ReserveForRun().
 	inst.reset(cancel)
-	p.mu.Unlock()
 
 	start := time.Now()
 	inst.status(StateBooting)
@@ -246,24 +243,25 @@ func (pi *poolInstance[T]) getInfo() Info {
 }
 
 func (pi *poolInstance[T]) reserve(ch chan Runner[T]) {
+	pi.mu.Lock()
+	// If we don't take the lock, it's possible that instance restart would race with job/jobChan update.
 	pi.stop()
 	pi.jobChan = ch
 	pi.job = nil
-	pi.updateInfo(func(info *Info) {
-		info.Reserved = true
-	})
+	pi.info.Reserved = true
+	pi.mu.Unlock()
 }
 
 func (pi *poolInstance[T]) free(job Runner[T]) {
+	pi.mu.Lock()
 	pi.job = job
 	pi.jobChan = nil
-
-	pi.updateInfo(func(info *Info) {
-		info.Reserved = false
-	})
+	switchToJob := pi.switchToJob
+	pi.info.Reserved = false
+	pi.mu.Unlock()
 
 	select {
-	case pi.switchToJob <- job:
+	case switchToJob <- job:
 		// Just in case the instance has been waiting.
 		return
 	default:
-- 
cgit mrf-deployment