fixes
This commit is contained in:
onysd
parent
21a0856587
commit
e8dc967e6a
26 changed files with 1373 additions and 481 deletions
|
|
@ -11,137 +11,257 @@ import (
|
|||
)
|
||||
|
||||
// BroadcastStore is the in-memory implementation of store.BroadcastStore,
|
||||
// used by admin/app unit tests.
|
||||
// used by admin/app unit tests. It has no concept of a "users table" to
|
||||
// snapshot against for "all" mode, so callers seed eligible user ids via
|
||||
// SeedEligibleUsers; MaterializeBroadcastRecipients walks that fixed set the
|
||||
// same way the postgres backend walks a keyset range.
|
||||
type BroadcastStore struct {
|
||||
mu sync.Mutex
|
||||
broadcasts map[int64]domain.Broadcast
|
||||
recipients map[int64]*memBroadcastRecipient
|
||||
nextBID int64
|
||||
nextRID int64
|
||||
}
|
||||
|
||||
type memBroadcastRecipient struct {
|
||||
domain.BroadcastRecipient
|
||||
message string
|
||||
mu sync.Mutex
|
||||
broadcasts map[int64]domain.Broadcast
|
||||
recipients map[int64]*domain.BroadcastRecipient
|
||||
eligibleUsers []int64 // sorted ascending, mirrors "all non-bot, non-system users"
|
||||
nextBID int64
|
||||
nextRID int64
|
||||
}
|
||||
|
||||
func NewBroadcastStore() *BroadcastStore {
|
||||
return &BroadcastStore{
|
||||
broadcasts: make(map[int64]domain.Broadcast),
|
||||
recipients: make(map[int64]*memBroadcastRecipient),
|
||||
recipients: make(map[int64]*domain.BroadcastRecipient),
|
||||
}
|
||||
}
|
||||
|
||||
var _ store.BroadcastStore = (*BroadcastStore)(nil)
|
||||
|
||||
func (s *BroadcastStore) CreateBroadcast(_ context.Context, message string, targetMode domain.BroadcastTargetMode, recipientUserIDs []int64, createdBy string) (domain.Broadcast, error) {
|
||||
if len(recipientUserIDs) == 0 {
|
||||
return domain.Broadcast{}, domain.ErrBroadcastNoRecipients
|
||||
}
|
||||
// SeedEligibleUsers sets the fixed set of user ids "all"-mode targets and
|
||||
// PreviewBroadcastRecipients/CreateBroadcast/MaterializeBroadcastRecipients
|
||||
// enumerate over, mirroring the postgres store's live users-table query.
|
||||
func (s *BroadcastStore) SeedEligibleUsers(userIDs []int64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.nextBID++
|
||||
b := domain.Broadcast{
|
||||
ID: s.nextBID,
|
||||
Message: message,
|
||||
TargetMode: targetMode,
|
||||
CreatedBy: createdBy,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
seen := make(map[int64]bool, len(recipientUserIDs))
|
||||
for _, userID := range recipientUserIDs {
|
||||
if seen[userID] {
|
||||
continue
|
||||
}
|
||||
seen[userID] = true
|
||||
s.nextRID++
|
||||
s.recipients[s.nextRID] = &memBroadcastRecipient{
|
||||
BroadcastRecipient: domain.BroadcastRecipient{
|
||||
ID: s.nextRID,
|
||||
BroadcastID: b.ID,
|
||||
UserID: userID,
|
||||
Status: domain.BroadcastRecipientPending,
|
||||
},
|
||||
message: message,
|
||||
}
|
||||
b.TotalCount++
|
||||
}
|
||||
s.broadcasts[b.ID] = b
|
||||
return b, nil
|
||||
s.eligibleUsers = append([]int64(nil), userIDs...)
|
||||
sort.Slice(s.eligibleUsers, func(i, j int) bool { return s.eligibleUsers[i] < s.eligibleUsers[j] })
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) PendingBroadcastRecipients(_ context.Context, limit int) ([]store.PendingBroadcastRecipient, error) {
|
||||
if limit <= 0 || limit > 200 {
|
||||
limit = 50
|
||||
func isEligibleSelected(userID int64) bool {
|
||||
return userID > 0 && !domain.IsSystemUserID(userID)
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) PreviewBroadcastRecipients(_ context.Context, mode domain.BroadcastTargetMode, selectedUserIDs []int64) (int64, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
switch mode {
|
||||
case domain.BroadcastTargetAll:
|
||||
if len(s.eligibleUsers) == 0 {
|
||||
return 0, domain.ErrBroadcastNoRecipients
|
||||
}
|
||||
return int64(len(s.eligibleUsers)), nil
|
||||
case domain.BroadcastTargetSelected:
|
||||
if len(selectedUserIDs) == 0 {
|
||||
return 0, domain.ErrBroadcastNoRecipients
|
||||
}
|
||||
for _, id := range selectedUserIDs {
|
||||
if !isEligibleSelected(id) {
|
||||
return 0, domain.ErrBroadcastRecipientInvalid
|
||||
}
|
||||
}
|
||||
return int64(len(selectedUserIDs)), nil
|
||||
default:
|
||||
return 0, domain.ErrBroadcastInvalid
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) CreateBroadcast(_ context.Context, message string, entities []domain.MessageEntity, mode domain.BroadcastTargetMode, selectedUserIDs []int64, createdBy string) (domain.Broadcast, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
switch mode {
|
||||
case domain.BroadcastTargetAll:
|
||||
if len(s.eligibleUsers) == 0 {
|
||||
return domain.Broadcast{}, domain.ErrBroadcastNoRecipients
|
||||
}
|
||||
s.nextBID++
|
||||
b := domain.Broadcast{
|
||||
ID: s.nextBID, Message: message, Entities: entities, TargetMode: mode,
|
||||
TargetCount: int64(len(s.eligibleUsers)), CreatedBy: createdBy, CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
s.broadcasts[b.ID] = b
|
||||
return b, nil
|
||||
case domain.BroadcastTargetSelected:
|
||||
if len(selectedUserIDs) == 0 {
|
||||
return domain.Broadcast{}, domain.ErrBroadcastNoRecipients
|
||||
}
|
||||
for _, id := range selectedUserIDs {
|
||||
if !isEligibleSelected(id) {
|
||||
return domain.Broadcast{}, domain.ErrBroadcastRecipientInvalid
|
||||
}
|
||||
}
|
||||
s.nextBID++
|
||||
b := domain.Broadcast{
|
||||
ID: s.nextBID, Message: message, Entities: entities, TargetMode: mode,
|
||||
EnumerationDone: true, CreatedBy: createdBy, CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
seen := make(map[int64]bool, len(selectedUserIDs))
|
||||
for _, userID := range selectedUserIDs {
|
||||
if seen[userID] {
|
||||
continue
|
||||
}
|
||||
seen[userID] = true
|
||||
s.nextRID++
|
||||
s.recipients[s.nextRID] = &domain.BroadcastRecipient{
|
||||
ID: s.nextRID, BroadcastID: b.ID, UserID: userID,
|
||||
Status: domain.BroadcastRecipientPending, NextAttemptAt: time.Now().UTC(),
|
||||
}
|
||||
b.TargetCount++
|
||||
b.MaterializedCount++
|
||||
}
|
||||
s.broadcasts[b.ID] = b
|
||||
return b, nil
|
||||
default:
|
||||
return domain.Broadcast{}, domain.ErrBroadcastInvalid
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) MaterializeBroadcastRecipients(_ context.Context, limit int) (int, error) {
|
||||
if limit <= 0 || limit > 1000 {
|
||||
limit = 100
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
// Iteration order over a map is unspecified; sort by recipient id (assigned
|
||||
// in creation order) so this matches the postgres backend's "oldest first".
|
||||
ids := make([]int64, 0, len(s.recipients))
|
||||
var ids []int64
|
||||
for id, b := range s.broadcasts {
|
||||
if b.TargetMode == domain.BroadcastTargetAll && !b.EnumerationDone {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
sortInt64s(ids)
|
||||
bid := ids[0]
|
||||
b := s.broadcasts[bid]
|
||||
inserted := 0
|
||||
for _, userID := range s.eligibleUsers {
|
||||
if int64(inserted) >= int64(limit) {
|
||||
break
|
||||
}
|
||||
if s.hasRecipient(bid, userID) {
|
||||
continue
|
||||
}
|
||||
s.nextRID++
|
||||
s.recipients[s.nextRID] = &domain.BroadcastRecipient{
|
||||
ID: s.nextRID, BroadcastID: bid, UserID: userID,
|
||||
Status: domain.BroadcastRecipientPending, NextAttemptAt: time.Now().UTC(),
|
||||
}
|
||||
b.MaterializedCount++
|
||||
inserted++
|
||||
}
|
||||
if inserted < limit {
|
||||
b.EnumerationDone = true
|
||||
b.TargetCount = b.MaterializedCount
|
||||
}
|
||||
s.broadcasts[bid] = b
|
||||
return inserted, nil
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) hasRecipient(broadcastID, userID int64) bool {
|
||||
for _, r := range s.recipients {
|
||||
if r.BroadcastID == broadcastID && r.UserID == userID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) ClaimBroadcastRecipients(_ context.Context, leaseToken string, limit int, lease time.Duration) ([]store.BroadcastRecipientClaim, error) {
|
||||
if leaseToken == "" {
|
||||
return nil, domain.ErrBroadcastInvalid
|
||||
}
|
||||
if limit <= 0 || limit > 500 {
|
||||
limit = 50
|
||||
}
|
||||
if lease <= 0 {
|
||||
lease = 30 * time.Second
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
var ids []int64
|
||||
now := time.Now().UTC()
|
||||
for id, r := range s.recipients {
|
||||
if r.Status == domain.BroadcastRecipientPending {
|
||||
eligible := (r.Status == domain.BroadcastRecipientPending && !r.NextAttemptAt.After(now)) ||
|
||||
(r.Status == domain.BroadcastRecipientProcessing && r.LeaseUntil != nil && !r.LeaseUntil.After(now))
|
||||
if eligible {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
sortInt64s(ids)
|
||||
out := make([]store.PendingBroadcastRecipient, 0, limit)
|
||||
if len(ids) > limit {
|
||||
ids = ids[:limit]
|
||||
}
|
||||
out := make([]store.BroadcastRecipientClaim, 0, len(ids))
|
||||
until := now.Add(lease)
|
||||
for _, id := range ids {
|
||||
if len(out) >= limit {
|
||||
break
|
||||
}
|
||||
r := s.recipients[id]
|
||||
out = append(out, store.PendingBroadcastRecipient{
|
||||
RecipientID: r.ID, BroadcastID: r.BroadcastID, UserID: r.UserID, Attempts: r.Attempts, Message: r.message,
|
||||
r.Status = domain.BroadcastRecipientProcessing
|
||||
r.Attempts++
|
||||
r.LeaseToken = leaseToken
|
||||
r.LeaseUntil = &until
|
||||
r.UpdatedAt = now
|
||||
b := s.broadcasts[r.BroadcastID]
|
||||
out = append(out, store.BroadcastRecipientClaim{
|
||||
RecipientID: r.ID, BroadcastID: r.BroadcastID, UserID: r.UserID,
|
||||
Attempts: r.Attempts, LeaseToken: leaseToken, Message: b.Message, Entities: b.Entities,
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) MarkBroadcastRecipientSent(_ context.Context, recipientID int64) error {
|
||||
func (s *BroadcastStore) CompleteBroadcastRecipient(_ context.Context, claim store.BroadcastRecipientClaim, privateMessageID int64, messageBoxID int, pts int) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
r, ok := s.recipients[recipientID]
|
||||
if !ok || r.Status != domain.BroadcastRecipientPending {
|
||||
return nil
|
||||
r, ok := s.recipients[claim.RecipientID]
|
||||
if !ok || r.Status != domain.BroadcastRecipientProcessing || r.LeaseToken != claim.LeaseToken {
|
||||
return domain.ErrBroadcastLeaseLost
|
||||
}
|
||||
r.Status = domain.BroadcastRecipientSent
|
||||
now := time.Now().UTC()
|
||||
r.SentAt = &now
|
||||
r.Status = domain.BroadcastRecipientSent
|
||||
r.LeaseToken = ""
|
||||
r.LeaseUntil = nil
|
||||
r.LastError = ""
|
||||
r.PrivateMessageID = privateMessageID
|
||||
r.MessageBoxID = messageBoxID
|
||||
r.Pts = pts
|
||||
r.SentAt = &now
|
||||
r.UpdatedAt = now
|
||||
b := s.broadcasts[claim.BroadcastID]
|
||||
b.SentCount++
|
||||
s.broadcasts[claim.BroadcastID] = b
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) MarkBroadcastRecipientFailed(_ context.Context, recipientID int64, reason string) error {
|
||||
func (s *BroadcastStore) ReleaseBroadcastRecipient(_ context.Context, claim store.BroadcastRecipientClaim, cause string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
r, ok := s.recipients[recipientID]
|
||||
if !ok || r.Status != domain.BroadcastRecipientPending {
|
||||
r, ok := s.recipients[claim.RecipientID]
|
||||
if !ok || r.Status != domain.BroadcastRecipientProcessing || r.LeaseToken != claim.LeaseToken {
|
||||
return nil
|
||||
}
|
||||
r.Attempts++
|
||||
r.LastError = reason
|
||||
now := time.Now().UTC()
|
||||
r.LeaseToken = ""
|
||||
r.LeaseUntil = nil
|
||||
r.LastError = cause
|
||||
r.UpdatedAt = now
|
||||
if r.Attempts >= domain.MaxBroadcastRecipientAttempts {
|
||||
r.Status = domain.BroadcastRecipientFailed
|
||||
b := s.broadcasts[claim.BroadcastID]
|
||||
b.FailedCount++
|
||||
s.broadcasts[claim.BroadcastID] = b
|
||||
} else {
|
||||
r.Status = domain.BroadcastRecipientPending
|
||||
r.NextAttemptAt = now
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) countsFor(broadcastID int64) (sent, failed int) {
|
||||
for _, r := range s.recipients {
|
||||
if r.BroadcastID != broadcastID {
|
||||
continue
|
||||
}
|
||||
switch r.Status {
|
||||
case domain.BroadcastRecipientSent:
|
||||
sent++
|
||||
case domain.BroadcastRecipientFailed:
|
||||
failed++
|
||||
}
|
||||
}
|
||||
return sent, failed
|
||||
}
|
||||
|
||||
func (s *BroadcastStore) ListBroadcasts(_ context.Context, beforeID int64, limit int) ([]domain.Broadcast, bool, error) {
|
||||
if limit <= 0 || limit > 200 {
|
||||
limit = 50
|
||||
|
|
@ -161,9 +281,7 @@ func (s *BroadcastStore) ListBroadcasts(_ context.Context, beforeID int64, limit
|
|||
}
|
||||
out := make([]domain.Broadcast, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
b := s.broadcasts[id]
|
||||
b.SentCount, b.FailedCount = s.countsFor(id)
|
||||
out = append(out, b)
|
||||
out = append(out, s.broadcasts[id])
|
||||
}
|
||||
return out, hasMore, nil
|
||||
}
|
||||
|
|
@ -175,7 +293,6 @@ func (s *BroadcastStore) BroadcastByID(_ context.Context, id int64) (domain.Broa
|
|||
if !ok {
|
||||
return domain.Broadcast{}, false, nil
|
||||
}
|
||||
b.SentCount, b.FailedCount = s.countsFor(id)
|
||||
return b, true, nil
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue