From b33bdde297a01dea534a7b72372a27984297b14b Mon Sep 17 00:00:00 2001 From: Copybara Service Date: Fri, 3 Jul 2026 09:57:55 -0700 Subject: [PATCH] Optimize unlimited broadcast performance SaveBroadcastMessage (called when a client is about to receive a broadcast message) tends to fail due to context cancelation (timeout) in large broadcasts. That's because broadcast messages to clients are processed sequentially using a mutex to prevent parallel writes to the task's broadcast allocation, which would lead to various issues depending on the db implementation. *Context: An allocation is a temporary, local quota of the total broadcast limit assigned to a specific task. On startup, tasks create an allocation of an active broadcast (usually 20% of the remaining messages), which allows it to send broadcast messages to clients without syncing with other tasks on how many messages are still remaining.* For broadcasts with no limit (`db.BroadcastUnlimited`), tasks generally don't need to care how many messages others have sent, and we can therefore skip the allocation mechanism and process messages in parallel. With this change: - Unlimited broadcasts (`Limit == db.BroadcastUnlimited`) skip database allocation creation. - Senders call `SaveBroadcastMessage` with an empty `AllocationID` for unlimited broadcasts. - Datastores (SQLite, MySQL, Spanner) skip updating `broadcast_allocations` if the `AllocationID` is empty. - The `Manager` skips acquiring the in-memory `i.lock` for unlimited broadcasts, allowing concurrent processing. - The `sent` count is tracked in-memory and flushed to `broadcasts.sent` when the manager is closed or the broadcast is cleaned up. PiperOrigin-RevId: 942179136 --- fleetspeak/src/server/db/store.go | 11 +- .../server/dbtesting/broadcaststore_suite.go | 104 +++++++++++++++++- .../server/internal/broadcasts/broadcasts.go | 62 ++++++----- fleetspeak/src/server/mysql/broadcaststore.go | 42 +++++-- .../src/server/spanner/broadcaststore.go | 47 ++++++-- .../src/server/sqlite/broadcaststore.go | 42 +++++-- fleetspeak/src/server/stats.go | 4 +- 7 files changed, 250 insertions(+), 62 deletions(-) diff --git a/fleetspeak/src/server/db/store.go b/fleetspeak/src/server/db/store.go index 03b528cf..6d3f073c 100644 --- a/fleetspeak/src/server/db/store.go +++ b/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. diff --git a/fleetspeak/src/server/dbtesting/broadcaststore_suite.go b/fleetspeak/src/server/dbtesting/broadcaststore_suite.go index 427bdb17..38368331 100644 --- a/fleetspeak/src/server/dbtesting/broadcaststore_suite.go +++ b/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, }) }) } diff --git a/fleetspeak/src/server/internal/broadcasts/broadcasts.go b/fleetspeak/src/server/internal/broadcasts/broadcasts.go index 47e2beb7..208a7ee4 100644 --- a/fleetspeak/src/server/internal/broadcasts/broadcasts.go +++ b/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)) } } diff --git a/fleetspeak/src/server/mysql/broadcaststore.go b/fleetspeak/src/server/mysql/broadcaststore.go index c4853b50..67bba13a 100644 --- a/fleetspeak/src/server/mysql/broadcaststore.go +++ b/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,7 +440,7 @@ 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()) @@ -430,6 +448,14 @@ func (d *Datastore) CleanupAllocation(ctx context.Context, bID ids.BroadcastID, 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 { diff --git a/fleetspeak/src/server/spanner/broadcaststore.go b/fleetspeak/src/server/spanner/broadcaststore.go index 5d4ed63c..442c7fcd 100644 --- a/fleetspeak/src/server/spanner/broadcaststore.go +++ b/fleetspeak/src/server/spanner/broadcaststore.go @@ -166,6 +166,27 @@ func (d *Datastore) trySaveBroadcastMessage(ctx context.Context, txn *spanner.Re return err } + sentCols := []string{"BroadcastID", "ClientID"} + ms := []*spanner.Mutation{spanner.InsertOrUpdate(d.broadcastSent, sentCols, []any{bid.Bytes(), cid.Bytes()})} + + if len(aid.Bytes()) == 0 { + row, err := txn.ReadRow(ctx, d.broadcasts, spanner.Key{bid.Bytes()}, []string{"MessageLimit"}) + if err != nil { + if spanner.ErrCode(err) == codes.NotFound { + return db.ErrBroadcastDisabled + } + return err + } + var limit int64 + if err := row.Columns(&limit); err != nil { + return err + } + if limit == 0 { + return db.ErrBroadcastDisabled + } + return txn.BufferWrite(ms) + } + row, err := txn.ReadRow(ctx, d.broadcastAllocations, spanner.Key{bid.Bytes(), aid.Bytes()}, []string{"Sent", "MessageLimit", "ExpiresTime"}) if err != nil { if spanner.ErrCode(err) == codes.NotFound { @@ -194,9 +215,7 @@ func (d *Datastore) trySaveBroadcastMessage(ctx context.Context, txn *spanner.Re } allocationCols := []string{"BroadcastID", "AllocationID", "Sent"} - sentCols := []string{"BroadcastID", "ClientID"} - ms := []*spanner.Mutation{spanner.Update(d.broadcastAllocations, allocationCols, []any{bid.Bytes(), aid.Bytes(), sent + 1})} - ms = append(ms, spanner.InsertOrUpdate(d.broadcastSent, sentCols, []any{bid.Bytes(), cid.Bytes()})) + ms = append(ms, spanner.Update(d.broadcastAllocations, allocationCols, []any{bid.Bytes(), aid.Bytes(), sent + 1})) return txn.BufferWrite(ms) } @@ -331,15 +350,14 @@ func (d *Datastore) tryCreateAllocation(ctx context.Context, txn *spanner.ReadWr } // CleanupAllocation implements db.BroadcastStore. -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 { _, err := d.dbClient.ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error { - return d.tryCleanupAllocation(ctx, txn, bid, aid) + return d.tryCleanupAllocation(ctx, txn, bid, aid, finalSent) }) return err } -func (d *Datastore) tryCleanupAllocation(ctx context.Context, txn *spanner.ReadWriteTransaction, bid ids.BroadcastID, aid ids.AllocationID) error { - +func (d *Datastore) tryCleanupAllocation(ctx context.Context, txn *spanner.ReadWriteTransaction, bid ids.BroadcastID, aid ids.AllocationID, finalSent uint64) error { row, err := txn.ReadRow(ctx, d.broadcasts, spanner.Key{bid.Bytes()}, []string{"Allocated", "Sent"}) if err != nil { return err @@ -349,15 +367,22 @@ func (d *Datastore) tryCleanupAllocation(ctx context.Context, txn *spanner.ReadW return err } - row, err = txn.ReadRow(ctx, d.broadcastAllocations, spanner.Key{bid.Bytes(), aid.Bytes()}, []string{"MessageLimit", "Sent"}) + if len(aid.Bytes()) == 0 { + // Unlimited broadcast: no allocation in DB, just update Sent count. + broadcastCols := []string{"BroadcastID", "Sent"} + ms := []*spanner.Mutation{spanner.Update(d.broadcasts, broadcastCols, []any{bid.Bytes(), bSent + int64(finalSent)})} + return txn.BufferWrite(ms) + } + + row, err = txn.ReadRow(ctx, d.broadcastAllocations, spanner.Key{bid.Bytes(), aid.Bytes()}, []string{"MessageLimit"}) if err != nil { if spanner.ErrCode(err) == codes.NotFound { return nil } return err } - var messageLimit, baSent int64 - if err := row.Columns(&messageLimit, &baSent); err != nil { + var messageLimit int64 + if err := row.Columns(&messageLimit); err != nil { return err } @@ -367,7 +392,7 @@ func (d *Datastore) tryCleanupAllocation(ctx context.Context, txn *spanner.ReadW } broadcastCols := []string{"BroadcastID", "Allocated", "Sent"} - ms := []*spanner.Mutation{spanner.Update(d.broadcasts, broadcastCols, []any{bid.Bytes(), int64(newAllocated), bSent + baSent})} + ms := []*spanner.Mutation{spanner.Update(d.broadcasts, broadcastCols, []any{bid.Bytes(), int64(newAllocated), bSent + int64(finalSent)})} ms = append(ms, spanner.Delete(d.broadcastAllocations, spanner.Key{bid.Bytes(), aid.Bytes()})) return txn.BufferWrite(ms) } diff --git a/fleetspeak/src/server/sqlite/broadcaststore.go b/fleetspeak/src/server/sqlite/broadcaststore.go index 7e02b9ba..9bcaf3c0 100644 --- a/fleetspeak/src/server/sqlite/broadcaststore.go +++ b/fleetspeak/src/server/sqlite/broadcaststore.go @@ -235,6 +235,28 @@ func (d *Datastore) SaveBroadcastMessage(ctx context.Context, msg *fspb.Message, } return d.runInTx(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.String()) + 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.String(), cID.String()); 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.String(), aID.String()) @@ -255,14 +277,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.String(), aID.String()); err != nil { + if _, err := tx.ExecContext(ctx, "INSERT INTO broadcast_sent(broadcast_id, client_id) VALUES (?, ?)", bID.String(), cID.String()); err != nil { return err } - _, err = tx.ExecContext(ctx, "INSERT INTO broadcast_sent(broadcast_id, client_id) VALUES (?, ?)", bID.String(), cID.String()) + _, err = tx.ExecContext(ctx, "UPDATE broadcast_allocations SET sent = ? WHERE broadcast_id = ? AND allocation_id = ?", as+1, bID.String(), aID.String()) return err }) } @@ -432,7 +450,7 @@ 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 { d.l.Lock() defer d.l.Unlock() return d.runInTx(func(tx *sql.Tx) error { @@ -442,6 +460,14 @@ func (d *Datastore) CleanupAllocation(ctx context.Context, bID ids.BroadcastID, 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.String()); 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.String(), aID.String()) if err := r.Scan(&as, &al); err != nil { @@ -454,7 +480,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.String()); err != nil { + if _, err := tx.ExecContext(ctx, "UPDATE broadcasts SET sent = ?, allocated = ? WHERE broadcast_id = ?", b.sent+finalSent, newAllocated, bID.String()); err != nil { return err } if _, err := tx.ExecContext(ctx, "DELETE from broadcast_allocations WHERE broadcast_id = ? AND allocation_id = ?", bID.String(), aID.String()); err != nil { diff --git a/fleetspeak/src/server/stats.go b/fleetspeak/src/server/stats.go index 65facf7e..dbf93e42 100644 --- a/fleetspeak/src/server/stats.go +++ b/fleetspeak/src/server/stats.go @@ -283,9 +283,9 @@ func (d MonitoredDatastore) CreateAllocation(ctx context.Context, id ids.Broadca return res, err } -func (d MonitoredDatastore) CleanupAllocation(ctx context.Context, bid ids.BroadcastID, aid ids.AllocationID) error { +func (d MonitoredDatastore) CleanupAllocation(ctx context.Context, bid ids.BroadcastID, aid ids.AllocationID, finalSent uint64) error { s := ftime.Now() - err := d.D.CleanupAllocation(ctx, bid, aid) + err := d.D.CleanupAllocation(ctx, bid, aid, finalSent) d.C.DatastoreOperation(s, ftime.Now(), "CleanupAllocation", err) return err }