Optimize unlimited broadcast performance

fleetspeak/src/server/db/store.go

@@ -360,6 +360,9 @@ type BroadcastStore interface {
 	DisableBroadcasts(ctx context.Context, bIDs []ids.BroadcastID) error
 
 	// SaveBroadcastMessage saves a new broadcast message.
+	// If aid is empty, saving the message does not update any allocation. If the
+	// broadcast identified by bid is disabled, SaveBroadcastMessage returns
+	// [ErrBroadcastDisabled].
 	SaveBroadcastMessage(ctx context.Context, msg *fspb.Message, bid ids.BroadcastID, cid common.ClientID, aid ids.AllocationID) error
 
 	// ListActiveBroadcasts lists broadcasts which could be sent to some
@@ -374,10 +377,10 @@ type BroadcastStore interface {
 	// expiry. Return nil if there is no message allocation available.
 	CreateAllocation(ctx context.Context, id ids.BroadcastID, frac float32, expiry time.Time) (*AllocationInfo, error)
 
-	// CleanupAllocation deletes the identified allocation record and
-	// updates the broadcast sent count according to the number that were
-	// actually sent under the given allocation.
-	CleanupAllocation(ctx context.Context, bid ids.BroadcastID, aid ids.AllocationID) error
+	// CleanupAllocation increases the broadcast sent count in the datastore
+	// according to finalSent and deletes the specified allocation record if aid
+	// is non-empty.
+	CleanupAllocation(ctx context.Context, bid ids.BroadcastID, aid ids.AllocationID, finalSent uint64) error
 }
 
 // ReadSeekerCloser groups io.ReadSeeker and io.Closer.

fleetspeak/src/server/dbtesting/broadcaststore_suite.go

@@ -184,7 +184,11 @@ func broadcastStoreTest(t *testing.T, ds db.Store) {
 
 	// Clean them all up.
 	for _, ids := range allocs {
-		if err := ds.CleanupAllocation(ctx, ids.bID, ids.aID); err != nil {
+		var finalSent uint64
+		if ids == allocs[0] || ids == allocs[4] {
+			finalSent = 1
+		}
+		if err := ds.CleanupAllocation(ctx, ids.bID, ids.aID, finalSent); err != nil {
 			t.Errorf("Unable to cleanup allocation %v: %v", ids, err)
 		}
 	}
@@ -248,8 +252,20 @@ func testDisableBroadcasts(t *testing.T, ds db.Store, bid ids.BroadcastID, clien
 		t.Errorf("Expected ErrBroadcastDisabled, got: %v", err)
 	}
 
+	// 3b. SaveBroadcastMessage with empty AllocationID (unlimited) should also fail with ErrBroadcastDisabled.
+	mid2, _ := common.RandomMessageID()
+	err = ds.SaveBroadcastMessage(ctx, &fspb.Message{
+		MessageId:    mid2.Bytes(),
+		Destination:  &fspb.Address{ClientId: clientID.Bytes(), ServiceName: "test"},
+		CreationTime: db.NowProto(),
+	}, bid, clientID, ids.AllocationID{})
+
+	if !errors.Is(err, db.ErrBroadcastDisabled) {
+		t.Errorf("Expected ErrBroadcastDisabled for empty AllocationID, got: %v", err)
+	}
+
 	// 4. CleanupAllocation should ignore the missing allocation and succeed gracefully.
-	if err := ds.CleanupAllocation(ctx, bid, a.ID); err != nil {
+	if err := ds.CleanupAllocation(ctx, bid, a.ID, 0); err != nil {
 		t.Errorf("CleanupAllocation failed for deleted allocation: %v", err)
 	}
 }
@@ -353,11 +369,95 @@ func testBroadcastReplacement(t *testing.T, ds db.Store) {
 	}
 }
 
+func testUnlimitedBroadcast(t *testing.T, ds db.Store) {
+	ctx := t.Context()
+
+	bID, err := ids.RandomBroadcastID()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	br := &spb.Broadcast{
+		BroadcastId: bID.Bytes(),
+		Source:      &fspb.Address{ServiceName: "testService"},
+		MessageType: "UnlimitedBroadcastNoAlloc",
+	}
+
+	if err := ds.CreateBroadcast(ctx, br, db.BroadcastUnlimited); err != nil {
+		t.Fatalf("CreateBroadcast failed: %v", err)
+	}
+
+	clientID, _ := common.BytesToClientID([]byte{0, 0, 0, 0, 0, 0, 0, 9})
+	if err := ds.AddClient(ctx, clientID, &db.ClientData{
+		Key: []byte("client key"),
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	// 1. Save broadcast message with empty AllocationID (unlimited broadcast bypass).
+	mID, _ := common.RandomMessageID()
+	msg := &fspb.Message{
+		MessageId:    mID.Bytes(),
+		Source:       &fspb.Address{ServiceName: "testService"},
+		Destination:  &fspb.Address{ClientId: clientID.Bytes(), ServiceName: "testService"},
+		MessageType:  "UnlimitedBroadcastNoAlloc",
+		CreationTime: db.NowProto(),
+	}
+
+	if err := ds.SaveBroadcastMessage(ctx, msg, bID, clientID, ids.AllocationID{}); err != nil {
+		t.Fatalf("SaveBroadcastMessage with empty AllocationID failed: %v", err)
+	}
+
+	// 2. Sent count in DB should still be 0 before CleanupAllocation.
+	bs, err := ds.ListActiveBroadcasts(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var found *db.BroadcastInfo
+	for _, b := range bs {
+		if bytes.Equal(b.Broadcast.BroadcastId, bID.Bytes()) {
+			found = b
+			break
+		}
+	}
+	if found == nil {
+		t.Fatalf("Broadcast %v not found in active broadcasts", bID)
+	}
+	if found.Sent != 0 {
+		t.Errorf("Expected Sent count in DB to be 0 before cleanup, got %v", found.Sent)
+	}
+
+	// 3. CleanupAllocation with empty AllocationID and finalSent = 1.
+	if err := ds.CleanupAllocation(ctx, bID, ids.AllocationID{}, 1); err != nil {
+		t.Fatalf("CleanupAllocation with empty AllocationID failed: %v", err)
+	}
+
+	// 4. Sent count in DB should now be updated to 1.
+	bs, err = ds.ListActiveBroadcasts(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	found = nil
+	for _, b := range bs {
+		if bytes.Equal(b.Broadcast.BroadcastId, bID.Bytes()) {
+			found = b
+			break
+		}
+	}
+	if found == nil {
+		t.Fatalf("Broadcast %v not found in active broadcasts after cleanup", bID)
+	}
+	if found.Sent != 1 {
+		t.Errorf("Expected Sent count in DB to be 1 after cleanup, got %v", found.Sent)
+	}
+}
+
 func broadcastStoreTestSuite(t *testing.T, env DbTestEnv) {
 	t.Run("BroadcastStoreTestSuite", func(t *testing.T) {
 		runTestSuite(t, env, map[string]func(*testing.T, db.Store){
 			"BroadcastStoreTest":       broadcastStoreTest,
 			"BroadcastReplacementTest": testBroadcastReplacement,
+			"UnlimitedBroadcastTest":   testUnlimitedBroadcast,
 		})
 	})
 }

fleetspeak/src/server/internal/broadcasts/broadcasts.go

@@ -158,14 +158,15 @@ Infos:
 				continue Infos
 			}
 		}
+		info.useCount.Add(1)
 		if info.limit == db.BroadcastUnlimited {
 			atomic.AddUint64(&info.sent, 1)
 		} else {
 			if !limitedAtomicIncrement(&info.sent, info.limit) {
+				info.useCount.Done()
 				continue
 			}
 		}
-		info.useCount.Add(1)
 		is = append(is, info)
 	}
 	m.l.RUnlock()
@@ -178,10 +179,8 @@ Infos:
 		mid, err := common.RandomMessageID()
 		if err != nil {
 			log.Errorf("unable to create message id: %v", err)
-			if i.limit != db.BroadcastUnlimited {
-				// Incantation to decrement a uint64, recommend AddUint64 docs:
-				atomic.AddUint64(&i.sent, ^uint64(0))
-			}
+			// Incantation to decrement a uint64, recommend AddUint64 docs:
+			atomic.AddUint64(&i.sent, ^uint64(0))
 			i.useCount.Done()
 			continue
 		}
@@ -196,17 +195,19 @@ Infos:
 			Data:         i.b.Data,
 			CreationTime: db.NowProto(),
 		}
-		i.lock.Lock()
-		err = m.bs.SaveBroadcastMessage(ctx, msg, i.bID, id, i.aID)
-		i.lock.Unlock()
+		if i.limit == db.BroadcastUnlimited {
+			err = m.bs.SaveBroadcastMessage(ctx, msg, i.bID, id, ids.AllocationID{})
+		} else {
+			i.lock.Lock()
+			err = m.bs.SaveBroadcastMessage(ctx, msg, i.bID, id, i.aID)
+			i.lock.Unlock()
+		}
 		if err != nil {
 			if !errors.Is(err, db.ErrBroadcastDisabled) {
 				log.Errorf("SaveBroadcastMessage of instance of broadcast %v failed. Not sending. [%v]", i.bID, err)
 			}
-			if i.limit != db.BroadcastUnlimited {
-				// Incantation to decrement a uint64, recommend by AddUint64 docs:
-				atomic.AddUint64(&i.sent, ^uint64(0))
-			}
+			// Incantation to decrement a uint64, recommend by AddUint64 docs:
+			atomic.AddUint64(&i.sent, ^uint64(0))
 			i.useCount.Done()
 			continue
 		}
@@ -260,22 +261,27 @@ func (m *Manager) refreshInfo(ctx context.Context) error {
 			if b.Sent == b.Limit {
 				continue
 			}
-			a, err := m.bs.CreateAllocation(ctx, id, allocFrac, ftime.Now().Add(allocDuration))
-			if err != nil {
-				log.Errorf("Unable to create alloc for broadcast %v, skipping: %v", id, err)
-				continue
+			info := &bInfo{
+				bID:   id,
+				b:     b.Broadcast,
+				limit: b.Limit,
+				sent:  0,
 			}
-			if a != nil {
-				newAllocs[id] = &bInfo{
-					bID: id,
-					b:   b.Broadcast,
-
-					aID:    a.ID,
-					limit:  a.Limit,
-					sent:   0,
-					expiry: a.Expiry,
+			if b.Limit != db.BroadcastUnlimited {
+				a, err := m.bs.CreateAllocation(ctx, id, allocFrac, ftime.Now().Add(allocDuration))
+				if err != nil {
+					log.Errorf("Unable to create alloc for broadcast %v, skipping: %v", id, err)
+					continue
+				}
+				if a == nil {
+					// No remaining messages to allocate.
+					continue
 				}
+				info.aID = a.ID
+				info.limit = a.Limit
+				info.expiry = a.Expiry
 			}
+			newAllocs[id] = info
 		}
 	}
 
@@ -309,7 +315,8 @@ func (m *Manager) refreshInfo(ctx context.Context) error {
 	// cleanup everything we can, even if there are errors.
 	for _, a := range c {
 		a.useCount.Wait()
-		if err := m.bs.CleanupAllocation(ctx, a.bID, a.aID); err != nil {
+		finalSent := atomic.LoadUint64(&a.sent)
+		if err := m.bs.CleanupAllocation(ctx, a.bID, a.aID, finalSent); err != nil {
 			errMsgs = append(errMsgs, fmt.Sprintf("[%v,%v]:\"%v\"", a.bID, a.aID, err))
 		}
 	}
@@ -378,7 +385,8 @@ func (m *Manager) Close(ctx context.Context) error {
 	var errMsgs []string
 	for _, i := range m.infos {
 		i.useCount.Wait()
-		if err := m.bs.CleanupAllocation(ctx, i.bID, i.aID); err != nil {
+		finalSent := atomic.LoadUint64(&i.sent)
+		if err := m.bs.CleanupAllocation(ctx, i.bID, i.aID, finalSent); err != nil {
 			errMsgs = append(errMsgs, fmt.Sprintf("[%v,%v]:\"%v\"", i.bID, i.aID, err))
 		}
 	}

fleetspeak/src/server/mysql/broadcaststore.go

@@ -229,6 +229,28 @@ func (d *Datastore) SaveBroadcastMessage(ctx context.Context, msg *fspb.Message,
 	}
 
 	return d.runInTx(ctx, false, func(tx *sql.Tx) error {
+		if err := d.tryStoreMessage(ctx, tx, dbm, true); err != nil {
+			return err
+		}
+
+		if len(aID.Bytes()) == 0 {
+			var limit uint64
+			r := tx.QueryRowContext(ctx, "SELECT message_limit FROM broadcasts WHERE broadcast_id = ?", bID.Bytes())
+			if err := r.Scan(&limit); err != nil {
+				if errors.Is(err, sql.ErrNoRows) {
+					return db.ErrBroadcastDisabled
+				}
+				return err
+			}
+			if limit == 0 {
+				return db.ErrBroadcastDisabled
+			}
+			if _, err := tx.ExecContext(ctx, "INSERT INTO broadcast_sent(broadcast_id, client_id) VALUES (?, ?)", bID.Bytes(), cID.Bytes()); err != nil {
+				return err
+			}
+			return nil
+		}
+
 		var as, al uint64
 		exp := &tspb.Timestamp{}
 		r := tx.QueryRowContext(ctx, "SELECT sent, message_limit, expiration_time_seconds, expiration_time_nanos FROM broadcast_allocations WHERE broadcast_id = ? AND allocation_id = ?", bID.Bytes(), aID.Bytes())
@@ -249,14 +271,10 @@ func (d *Datastore) SaveBroadcastMessage(ctx context.Context, msg *fspb.Message,
 			return fmt.Errorf("SaveBroadcastMessage: broadcast allocation [%v, %v] is expired: %v", aID, bID, et)
 		}
 
-		if err := d.tryStoreMessage(ctx, tx, dbm, true); err != nil {
-			return err
-		}
-
-		if _, err := tx.ExecContext(ctx, "UPDATE broadcast_allocations SET sent = ? WHERE broadcast_id = ? AND allocation_id = ?", as+1, bID.Bytes(), aID.Bytes()); err != nil {
+		if _, err := tx.ExecContext(ctx, "INSERT INTO broadcast_sent(broadcast_id, client_id) VALUES (?, ?)", bID.Bytes(), cID.Bytes()); err != nil {
 			return err
 		}
-		_, err = tx.ExecContext(ctx, "INSERT INTO broadcast_sent(broadcast_id, client_id) VALUES (?, ?)", bID.Bytes(), cID.Bytes())
+		_, err = tx.ExecContext(ctx, "UPDATE broadcast_allocations SET sent = ? WHERE broadcast_id = ? AND allocation_id = ?", as+1, bID.Bytes(), aID.Bytes())
 		return err
 	})
 }
@@ -422,14 +440,22 @@ func (d *Datastore) CreateAllocation(ctx context.Context, id ids.BroadcastID, fr
 	return ret, err
 }
 
-func (d *Datastore) CleanupAllocation(ctx context.Context, bID ids.BroadcastID, aID ids.AllocationID) error {
+func (d *Datastore) CleanupAllocation(ctx context.Context, bID ids.BroadcastID, aID ids.AllocationID, finalSent uint64) error {
 	return d.runInTx(ctx, false, func(tx *sql.Tx) error {
 		var b dbBroadcast
 		r := tx.QueryRowContext(ctx, "SELECT sent, allocated, message_limit FROM broadcasts WHERE broadcast_id = ?", bID.Bytes())
 		if err := r.Scan(&b.sent, &b.allocated, &b.messageLimit); err != nil {
 			return err
 		}
 
+		if len(aID.Bytes()) == 0 {
+			// Unlimited broadcast: no allocation in DB, just update Sent count.
+			if _, err := tx.ExecContext(ctx, "UPDATE broadcasts SET sent = ? WHERE broadcast_id = ?", b.sent+finalSent, bID.Bytes()); err != nil {
+				return err
+			}
+			return nil
+		}
+
 		var as, al uint64
 		r = tx.QueryRowContext(ctx, "SELECT sent, message_limit FROM broadcast_allocations WHERE broadcast_id = ? AND allocation_id = ?", bID.Bytes(), aID.Bytes())
 		if err := r.Scan(&as, &al); err != nil {
@@ -442,7 +468,7 @@ func (d *Datastore) CleanupAllocation(ctx context.Context, bID ids.BroadcastID,
 		if err != nil {
 			return fmt.Errorf("unable to clear allocation [%v,%v]: %v", bID.String(), aID.String(), err)
 		}
-		if _, err := tx.ExecContext(ctx, "UPDATE broadcasts SET sent = ?, allocated = ? WHERE broadcast_id = ?", b.sent+as, newAllocated, bID.Bytes()); err != nil {
+		if _, err := tx.ExecContext(ctx, "UPDATE broadcasts SET sent = ?, allocated = ? WHERE broadcast_id = ?", b.sent+finalSent, newAllocated, bID.Bytes()); err != nil {
 			return err
 		}
 		if _, err := tx.ExecContext(ctx, "DELETE from broadcast_allocations WHERE broadcast_id = ? AND allocation_id = ?", bID.Bytes(), aID.Bytes()); err != nil {