Bugfix: create/update manual label returns outdated info when configured with a mysql replica (#23725)

changes/23597-fix-create-update-label-returns-outdated-info

@@ -0,0 +1 @@
+* Fixed a bug where the create and update label endpoints could return outdated information in a deployment using a mysql replica.

server/datastore/mysql/testing_utils.go

@@ -110,12 +110,17 @@ func setupDummyReplica(t testing.TB, testName string, ds *Datastore, opts *Datas
 	}
 	t.Cleanup(cancel)
 
+	type replicationRun struct {
+		forceTables     []string
+		replicationDone chan struct{}
+	}
+
 	// start the replication goroutine that runs when signalled through a
 	// channel, the replication runs in lock-step - the test is in control of
 	// when the replication happens, by calling opts.RunReplication(), and when
 	// that call returns, the replication is guaranteed to be done. This supports
 	// simulating all kinds of replica lag.
-	ch := make(chan chan struct{})
+	ch := make(chan replicationRun)
 	go func() {
 		// if it exits because of a panic/failed replication, cancel the context
 		// immediately so that RunReplication is unblocked too.
@@ -168,6 +173,20 @@ func setupDummyReplica(t testing.TB, testName string, ds *Datastore, opts *Datas
             update_time >= ?`, testName, last)
 				require.NoError(t, err)
 
+				// dedupe and add forced tables
+				tableSet := make(map[string]bool, len(tables)+len(out.forceTables))
+				for _, tbl := range tables {
+					tableSet[tbl] = true
+				}
+				for _, tbl := range out.forceTables {
+					tableSet[tbl] = true
+				}
+				tables = tables[:0]
+				for tbl := range tableSet {
+					tables = append(tables, tbl)
+				}
+				t.Logf("changed tables since %v: %v", last, tables)
+
 				err = primary.GetContext(ctx, &last, `
           SELECT
             MAX(update_time)
@@ -177,6 +196,7 @@ func setupDummyReplica(t testing.TB, testName string, ds *Datastore, opts *Datas
             table_schema = ? AND
             table_type = 'BASE TABLE'`, testName)
 				require.NoError(t, err)
+				t.Logf("last update time of primary is now %v", last)
 
 				// replicate by dropping the existing table and re-creating it from
 				// the primary.
@@ -195,7 +215,7 @@ func setupDummyReplica(t testing.TB, testName string, ds *Datastore, opts *Datas
 					require.NoError(t, err)
 				}
 
-				out <- struct{}{}
+				out.replicationDone <- struct{}{}
 				t.Logf("replication step executed, next will consider updates since %s", last)
 
 			case <-ctx.Done():
@@ -206,9 +226,9 @@ func setupDummyReplica(t testing.TB, testName string, ds *Datastore, opts *Datas
 
 	// set RunReplication to a function that triggers the replication and waits
 	// for it to complete.
-	opts.RunReplication = func() {
+	opts.RunReplication = func(forceTables ...string) {
 		done := make(chan struct{})
-		ch <- done
+		ch <- replicationRun{forceTables, done}
 		select {
 		case <-done:
 		case <-ctx.Done():
@@ -405,7 +425,10 @@ type DatastoreTestOptions struct {
 	// missing changes from the primary to the replica. The function is created
 	// and set automatically by CreateMySQLDSWithOptions. The test is in full
 	// control of when the replication is executed. Only applies to DummyReplica.
-	RunReplication func()
+	// Note that not all changes to data show up in the information_schema
+	// update_time timestamp, so to work around that limitation, explicit table
+	// names can be provided to force their replication.
+	RunReplication func(forceTables ...string)
 
 	// RealReplica indicates that the replica should be a real DB replica, with a dedicated connection.
 	RealReplica bool

server/service/labels.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 
+	"github.com/fleetdm/fleet/v4/server/contexts/ctxdb"
 	"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
 	"github.com/fleetdm/fleet/v4/server/contexts/license"
 	"github.com/fleetdm/fleet/v4/server/contexts/viewer"
@@ -109,6 +110,7 @@ func (svc *Service) NewLabel(ctx context.Context, p fleet.LabelPayload) (*fleet.
 		}
 
 		// must reload it to get the host IDs, refresh its count
+		ctx = ctxdb.RequirePrimary(ctx, true)
 		label, hostIDs, err = svc.ds.Label(ctx, label.ID, filter)
 		if err != nil {
 			return nil, nil, err
@@ -211,6 +213,7 @@ func (svc *Service) ModifyLabel(ctx context.Context, id uint, payload fleet.Modi
 			return svc.ds.SaveLabel(ctx, label, filter)
 		}
 		// Otherwise, simply reload label to get the host counts information
+		ctx = ctxdb.RequirePrimary(ctx, true)
 		return svc.ds.Label(ctx, id, filter)
 	}
 	return svc.ds.SaveLabel(ctx, label, filter)

server/service/labels_test.go

@@ -3,6 +3,7 @@ package service
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/fleetdm/fleet/v4/server/contexts/viewer"
 	"github.com/fleetdm/fleet/v4/server/datastore/mysql"
@@ -10,6 +11,7 @@ import (
 	"github.com/fleetdm/fleet/v4/server/mock"
 	"github.com/fleetdm/fleet/v4/server/ptr"
 	"github.com/fleetdm/fleet/v4/server/test"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -251,5 +253,62 @@ func TestApplyLabelSpecsWithBuiltInLabels(t *testing.T) {
 	}
 	err = svc.ApplyLabelSpecs(ctx, []*fleet.LabelSpec{spec})
 	assert.ErrorIs(t, err, assert.AnError)
+}
 
+func TestLabelsWithReplica(t *testing.T) {
+	opts := &mysql.DatastoreTestOptions{DummyReplica: true}
+	ds := mysql.CreateMySQLDSWithOptions(t, opts)
+	defer ds.Close()
+
+	svc, ctx := newTestService(t, ds, nil, nil)
+	ctx = viewer.NewContext(ctx, viewer.Viewer{User: &fleet.User{GlobalRole: ptr.String(fleet.RoleAdmin)}})
+
+	// create a couple hosts
+	h1, err := ds.NewHost(ctx, &fleet.Host{
+		Hostname:        "host1",
+		HardwareSerial:  uuid.NewString(),
+		UUID:            uuid.NewString(),
+		Platform:        "darwin",
+		LastEnrolledAt:  time.Now(),
+		DetailUpdatedAt: time.Now(),
+	})
+	require.NoError(t, err)
+	h2, err := ds.NewHost(ctx, &fleet.Host{
+		Hostname:        "host2",
+		HardwareSerial:  uuid.NewString(),
+		UUID:            uuid.NewString(),
+		Platform:        "darwin",
+		LastEnrolledAt:  time.Now(),
+		DetailUpdatedAt: time.Now(),
+	})
+	require.NoError(t, err)
+	// make the newly-created hosts available to the reader
+	opts.RunReplication()
+
+	lbl, hostIDs, err := svc.NewLabel(ctx, fleet.LabelPayload{Name: "label1", Hosts: []string{"host1", "host2"}})
+	require.NoError(t, err)
+	require.ElementsMatch(t, []uint{h1.ID, h2.ID}, hostIDs)
+	require.Equal(t, 2, lbl.HostCount)
+
+	// make the newly-created label available to the reader
+	opts.RunReplication("labels", "label_membership")
+
+	lbl, hostIDs, err = svc.ModifyLabel(ctx, lbl.ID, fleet.ModifyLabelPayload{Hosts: []string{"host1"}})
+	require.NoError(t, err)
+	require.ElementsMatch(t, []uint{h1.ID}, hostIDs)
+	require.Equal(t, 1, lbl.HostCount)
+
+	// reading this label without replication returns the old data as it only uses the reader
+	lbl, hostIDs, err = svc.GetLabel(ctx, lbl.ID)
+	require.NoError(t, err)
+	require.ElementsMatch(t, []uint{h1.ID, h2.ID}, hostIDs)
+	require.Equal(t, 2, lbl.HostCount)
+
+	// running the replication makes the updated data available
+	opts.RunReplication("labels", "label_membership")
+
+	lbl, hostIDs, err = svc.GetLabel(ctx, lbl.ID)
+	require.NoError(t, err)
+	require.ElementsMatch(t, []uint{h1.ID}, hostIDs)
+	require.Equal(t, 1, lbl.HostCount)
 }