From b1f6d8b07137aead6db06fa3defb10a0530884c9 Mon Sep 17 00:00:00 2001
From: matst80 <mats.tornberg@gmail.com>
Date: Sat, 9 Nov 2024 22:06:30 +0100
Subject: [PATCH] more mutex

---
 rpc-pool.go    | 6 ++++++
 synced-pool.go | 7 +++++++
 2 files changed, 13 insertions(+)

diff --git a/rpc-pool.go b/rpc-pool.go
index 6d1e16e..b7012b7 100644
--- a/rpc-pool.go
+++ b/rpc-pool.go
@@ -5,10 +5,12 @@ import (
 	"io"
 	"net"
 	"strings"
+	"sync"
 	"time"
 )
 
 type RemoteGrainPool struct {
+	mu     sync.RWMutex
 	Host   string
 	grains map[CartId]RemoteGrain
 }
@@ -87,7 +89,9 @@ func NewRemoteGrainPool(addr string) *RemoteGrainPool {
 }
 
 func (p *RemoteGrainPool) findRemoteGrain(id CartId) *RemoteGrain {
+	p.mu.RLock()
 	grain, ok := p.grains[id]
+	p.mu.RUnlock()
 	if !ok {
 		return nil
 	}
@@ -99,7 +103,9 @@ func (p *RemoteGrainPool) findOrCreateGrain(id CartId) *RemoteGrain {
 	grain := p.findRemoteGrain(id)
 	if grain == nil {
 		grain = NewRemoteGrain(id, p.Host)
+		p.mu.Lock()
 		p.grains[id] = *grain
+		p.mu.Unlock()
 		grain.Connect()
 	}
 	return grain
diff --git a/synced-pool.go b/synced-pool.go
index 0def54b..c4e8398 100644
--- a/synced-pool.go
+++ b/synced-pool.go
@@ -28,6 +28,7 @@ type RemoteHost struct {
 }
 
 type SyncedPool struct {
+	mu          sync.RWMutex
 	Discovery   Discovery
 	listener    net.Listener
 	Hostname    string
@@ -296,7 +297,9 @@ func (p *SyncedPool) handleConnection(conn net.Conn) {
 				if r.Host == string(idAndHostParts[1]) {
 					found = true
 					log.Printf("Remote grain %s changed to %s\n", idAndHostParts[0], idAndHostParts[1])
+					p.mu.Lock()
 					p.remoteIndex[ToCartId(idAndHostParts[0])] = r.Pool
+					p.mu.Unlock()
 				}
 			}
 
@@ -446,7 +449,9 @@ func (p *SyncedPool) Process(id CartId, messages ...Message) ([]byte, error) {
 	_, ok := p.local.grains[id]
 	if !ok {
 		// check if remote grain exists
+		p.mu.RLock()
 		remoteGrain, ok := p.remoteIndex[id]
+		p.mu.RUnlock()
 		if ok {
 			remoteLookupCount.Inc()
 			return remoteGrain.Process(id, messages...)
@@ -465,7 +470,9 @@ func (p *SyncedPool) Get(id CartId) ([]byte, error) {
 	_, ok := p.local.grains[id]
 	if !ok {
 		// check if remote grain exists
+		p.mu.RLock()
 		remoteGrain, ok := p.remoteIndex[id]
+		p.mu.RUnlock()
 		if ok {
 			remoteLookupCount.Inc()
 			return remoteGrain.Get(id)