introduce rw lock for db

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
juanfont · Jul 17, 2023 · 55aca05 · 55aca05
1 parent d38aca0
commit 55aca05
Show file tree

Hide file tree

Showing 7 changed files with 208 additions and 4 deletions.
diff --git a/hscontrol/db/addresses.go b/hscontrol/db/addresses.go
@@ -18,6 +18,9 @@ import (
 var ErrCouldNotAllocateIP = errors.New("could not find any suitable IP")
 
 func (hsdb *HSDatabase) getAvailableIPs() (types.MachineAddresses, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	var ips types.MachineAddresses
 	var err error
 	for _, ipPrefix := range hsdb.ipPrefixes {
@@ -33,6 +36,9 @@ func (hsdb *HSDatabase) getAvailableIPs() (types.MachineAddresses, error) {
 }
 
 func (hsdb *HSDatabase) getAvailableIP(ipPrefix netip.Prefix) (*netip.Addr, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	usedIps, err := hsdb.getUsedIPs()
 	if err != nil {
 		return nil, err

diff --git a/hscontrol/db/api_key.go b/hscontrol/db/api_key.go
@@ -22,6 +22,9 @@ var ErrAPIKeyFailedToParse = errors.New("failed to parse ApiKey")
 func (hsdb *HSDatabase) CreateAPIKey(
 	expiration *time.Time,
 ) (string, *types.APIKey, error) {
+	hsdb.mu.Lock()
+	defer hsdb.mu.Unlock()
+
 	prefix, err := util.GenerateRandomStringURLSafe(apiPrefixLength)
 	if err != nil {
 		return "", nil, err
@@ -55,6 +58,9 @@ func (hsdb *HSDatabase) CreateAPIKey(
 
 // ListAPIKeys returns the list of ApiKeys for a user.
 func (hsdb *HSDatabase) ListAPIKeys() ([]types.APIKey, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	keys := []types.APIKey{}
 	if err := hsdb.db.Find(&keys).Error; err != nil {
 		return nil, err
@@ -65,6 +71,9 @@ func (hsdb *HSDatabase) ListAPIKeys() ([]types.APIKey, error) {
 
 // GetAPIKey returns a ApiKey for a given key.
 func (hsdb *HSDatabase) GetAPIKey(prefix string) (*types.APIKey, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	key := types.APIKey{}
 	if result := hsdb.db.First(&key, "prefix = ?", prefix); result.Error != nil {
 		return nil, result.Error
@@ -75,6 +84,9 @@ func (hsdb *HSDatabase) GetAPIKey(prefix string) (*types.APIKey, error) {
 
 // GetAPIKeyByID returns a ApiKey for a given id.
 func (hsdb *HSDatabase) GetAPIKeyByID(id uint64) (*types.APIKey, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	key := types.APIKey{}
 	if result := hsdb.db.Find(&types.APIKey{ID: id}).First(&key); result.Error != nil {
 		return nil, result.Error
@@ -86,6 +98,9 @@ func (hsdb *HSDatabase) GetAPIKeyByID(id uint64) (*types.APIKey, error) {
 // DestroyAPIKey destroys a ApiKey. Returns error if the ApiKey
 // does not exist.
 func (hsdb *HSDatabase) DestroyAPIKey(key types.APIKey) error {
+	hsdb.mu.Lock()
+	defer hsdb.mu.Unlock()
+
 	if result := hsdb.db.Unscoped().Delete(key); result.Error != nil {
 		return result.Error
 	}
@@ -95,6 +110,9 @@ func (hsdb *HSDatabase) DestroyAPIKey(key types.APIKey) error {
 
 // ExpireAPIKey marks a ApiKey as expired.
 func (hsdb *HSDatabase) ExpireAPIKey(key *types.APIKey) error {
+	hsdb.mu.Lock()
+	defer hsdb.mu.Unlock()
+
 	if err := hsdb.db.Model(&key).Update("Expiration", time.Now()).Error; err != nil {
 		return err
 	}
@@ -103,6 +121,9 @@ func (hsdb *HSDatabase) ExpireAPIKey(key *types.APIKey) error {
 }
 
 func (hsdb *HSDatabase) ValidateAPIKey(keyStr string) (bool, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	prefix, hash, found := strings.Cut(keyStr, ".")
 	if !found {
 		return false, ErrAPIKeyFailedToParse

diff --git a/hscontrol/db/db.go b/hscontrol/db/db.go
@@ -40,6 +40,8 @@ type HSDatabase struct {
 	db       *gorm.DB
 	notifier *notifier.Notifier
 
+	mu sync.RWMutex
+
 	ipAllocationMutex sync.Mutex
 
 	ipPrefixes []netip.Prefix