Fix lost-update race in channel subscriptions

server/subscription_race_test.go

@@ -0,0 +1,114 @@
+// Copyright (c) 2019-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/mattermost/mattermost/server/public/model"
+	"github.com/mattermost/mattermost/server/public/plugin/plugintest"
+	"github.com/mattermost/mattermost/server/public/pluginapi"
+
+	"github.com/mattermost/mattermost-plugin-gitlab/server/subscription"
+)
+
+// raceKVAPI is an in-memory plugin.API whose KVGet and KVSetWithOptions
+// implement real compare-and-set semantics, so SetAtomicWithRetries behaves
+// exactly as it does against the server KV store. Every other API method is
+// inherited from plugintest.API and is unused by the subscription store.
+type raceKVAPI struct {
+	plugintest.API
+	mu    sync.Mutex
+	store map[string][]byte
+}
+
+func newRaceKVAPI() *raceKVAPI {
+	return &raceKVAPI{store: map[string][]byte{}}
+}
+
+func (a *raceKVAPI) KVGet(key string) ([]byte, *model.AppError) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	return a.store[key], nil
+}
+
+func (a *raceKVAPI) KVSetWithOptions(key string, value []byte, options model.PluginKVSetOptions) (bool, *model.AppError) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	if options.Atomic {
+		if !bytes.Equal(a.store[key], options.OldValue) {
+			return false, nil
+		}
+	}
+
+	if value == nil {
+		delete(a.store, key)
+	} else {
+		a.store[key] = value
+	}
+	return true, nil
+}
+
+// TestSubscriptionRace drives the plugin's real AddSubscription against an
+// in-memory KV store with real compare-and-set semantics. Many channels
+// subscribe to the same repository concurrently; afterwards we compare how many
+// AddSubscription calls reported success (nil error) against how many
+// subscriptions actually persisted.
+//
+// Correct behaviour: persisted == reported-success (every reported success is
+// durable). Before the fix, StoreSubscriptions did a plain non-atomic Set and
+// swallowed the error, so concurrent writers clobbered each other and persisted
+// was far lower than reported-success (silent lost updates).
+func TestSubscriptionRace(t *testing.T) {
+	const numChannels = 200
+
+	api := newRaceKVAPI()
+	p := &Plugin{configuration: &configuration{}}
+	p.SetAPI(api)
+	p.client = pluginapi.NewClient(api, p.Driver)
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	reportedSuccess := 0
+
+	for i := 0; i < numChannels; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			sub := &subscription.Subscription{
+				ChannelID:  fmt.Sprintf("channel-%d", i),
+				Repository: "namespace/project",
+			}
+			if _, err := p.AddSubscription("namespace/project", sub); err == nil {
+				mu.Lock()
+				reportedSuccess++
+				mu.Unlock()
+			}
+		}(i)
+	}
+	wg.Wait()
+
+	subs, err := p.GetSubscriptions()
+	if err != nil {
+		t.Fatalf("GetSubscriptions failed: %v", err)
+	}
+	persisted := 0
+	for _, chans := range subs.Repositories {
+		persisted += len(chans)
+	}
+
+	silentlyLost := reportedSuccess - persisted
+	t.Logf("concurrent AddSubscription calls : %d", numChannels)
+	t.Logf("AddSubscription returned success : %d", reportedSuccess)
+	t.Logf("subscriptions actually persisted : %d", persisted)
+	t.Logf("silently lost (success but gone) : %d", silentlyLost)
+
+	if silentlyLost > 0 {
+		t.Fatalf("lost-update bug: %d subscriptions were reported as saved but silently dropped", silentlyLost)
+	}
+}

server/subscriptions.go

@@ -5,6 +5,7 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 
 	"github.com/pkg/errors"
 
@@ -63,31 +64,28 @@ func filterSubscriptionsByChannel(subs *Subscriptions, channelID string) []*subs
 }
 
 func (p *Plugin) AddSubscription(fullPath string, sub *subscription.Subscription) (*Subscriptions, error) {
-	subs, err := p.GetSubscriptions()
-	if err != nil {
-		return nil, err
-	}
-
-	repoSubs := subs.Repositories[fullPath]
-	if repoSubs == nil {
-		repoSubs = []*subscription.Subscription{sub}
-	} else {
-		exists := false
-		for index, s := range repoSubs {
-			if s.ChannelID == sub.ChannelID {
-				repoSubs[index] = sub
-				exists = true
-				break
+	return p.modifySubscriptions(func(subs *Subscriptions) error {
+		repoSubs := subs.Repositories[fullPath]
+		if repoSubs == nil {
+			repoSubs = []*subscription.Subscription{sub}
+		} else {
+			exists := false
+			for index, s := range repoSubs {
+				if s.ChannelID == sub.ChannelID {
+					repoSubs[index] = sub
+					exists = true
+					break
+				}
 			}
-		}
 
-		if !exists {
-			repoSubs = append(repoSubs, sub)
+			if !exists {
+				repoSubs = append(repoSubs, sub)
+			}
 		}
-	}
 
-	subs.Repositories[fullPath] = repoSubs
-	return subs, p.StoreSubscriptions(subs)
+		subs.Repositories[fullPath] = repoSubs
+		return nil
+	})
 }
 
 func (p *Plugin) GetSubscriptions() (*Subscriptions, error) {
@@ -106,11 +104,44 @@ func (p *Plugin) GetSubscriptions() (*Subscriptions, error) {
 	return subscriptions, nil
 }
 
-func (p *Plugin) StoreSubscriptions(s *Subscriptions) error {
-	if _, err := p.client.KV.Set(SubscriptionsKey, s); err != nil {
-		p.client.Log.Warn("can't set subscriptions in kvstore", "err", err.Error())
+// errStopModify is a sentinel returned from a modifySubscriptions mutate
+// function to abort the atomic write without performing it and without treating
+// it as a store failure. SetAtomicWithRetries returns immediately (no retry)
+// when the callback errors, so the caller distinguishes this from a real error.
+var errStopModify = errors.New("stop modifying subscriptions")
+
+// modifySubscriptions performs an atomic read-modify-write of the whole
+// subscriptions blob. The mutate function runs inside SetAtomicWithRetries'
+// callback, so on every retry it receives the freshly re-read state and
+// re-applies its change on top of it. This prevents concurrent mutations across
+// channels from silently clobbering one another (lost updates). It returns the
+// resulting subscriptions on success.
+func (p *Plugin) modifySubscriptions(mutate func(*Subscriptions) error) (*Subscriptions, error) {
+	var result *Subscriptions
+
+	err := p.client.KV.SetAtomicWithRetries(SubscriptionsKey, func(oldValue []byte) (any, error) {
+		subs := &Subscriptions{Repositories: map[string][]*subscription.Subscription{}}
+		if len(oldValue) > 0 {
+			if err := json.Unmarshal(oldValue, subs); err != nil {
+				return nil, errors.Wrap(err, "can't unmarshal subscriptions from kvstore")
+			}
+			if subs.Repositories == nil {
+				subs.Repositories = map[string][]*subscription.Subscription{}
+			}
+		}
+
+		if err := mutate(subs); err != nil {
+			return nil, err
+		}
+
+		result = subs
+		return subs, nil
+	})
+	if err != nil {
+		return nil, err
 	}
-	return nil
+
+	return result, nil
 }
 
 func (p *Plugin) GetSubscribedChannelsForProject(
@@ -161,41 +192,55 @@ func (p *Plugin) Unsubscribe(channelID string, fullPath string) (bool, *Subscrip
 		return false, nil, errors.New("invalid repository")
 	}
 
-	subs, err := p.GetSubscriptions()
-	if err != nil {
-		return false, nil, err
-	}
-
 	var removed bool
+	var current *Subscriptions
 
-	// We don't know whether fullPath is a namespace or project, so we have to check both cases
-	for _, path := range []string{fullPath, fullPath + "/"} {
-		pathSubs := subs.Repositories[path]
-		if pathSubs == nil {
-			continue
-		}
+	subs, err := p.modifySubscriptions(func(subs *Subscriptions) error {
+		removed = false
+		current = subs
 
-		pathRemoved := false
-		for index, sub := range pathSubs {
-			if sub.ChannelID == channelID {
-				pathSubs = append(pathSubs[:index], pathSubs[index+1:]...)
-				pathRemoved = true
-				break
+		// We don't know whether fullPath is a namespace or project, so we have to check both cases
+		for _, path := range []string{fullPath, fullPath + "/"} {
+			pathSubs := subs.Repositories[path]
+			if pathSubs == nil {
+				continue
+			}
+
+			pathRemoved := false
+			for index, sub := range pathSubs {
+				if sub.ChannelID == channelID {
+					pathSubs = append(pathSubs[:index], pathSubs[index+1:]...)
+					pathRemoved = true
+					break
+				}
 			}
-		}
 
-		if pathRemoved {
-			if len(pathSubs) > 0 {
-				subs.Repositories[path] = pathSubs
-			} else {
-				delete(subs.Repositories, path)
+			if pathRemoved {
+				if len(pathSubs) > 0 {
+					subs.Repositories[path] = pathSubs
+				} else {
+					delete(subs.Repositories, path)
+				}
+				removed = true
 			}
-			removed = true
 		}
-	}
 
-	if !removed {
-		return false, subs, nil
+		if !removed {
+			return errStopModify
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		// errStopModify signals that no matching subscription existed; that is a
+		// no-op, not a failure. Any other error is a real store/read failure and
+		// must be surfaced instead of being reported as "not subscribed".
+		if errors.Cause(err) == errStopModify {
+			return false, current, nil
+		}
+		return false, nil, err
 	}
-	return true, subs, p.StoreSubscriptions(subs)
+
+	return true, subs, nil
 }