diff --git a/topic/service.go b/topic/service.go new file mode 100644 index 00000000..5fd89604 --- /dev/null +++ b/topic/service.go @@ -0,0 +1,521 @@ +// Package topic keeps the cluster-wide state of a topic that every node needs but no node owns: +// when it was first seen, when a subscriber was last active, and who pays for the pushes to it +// (its rate visitor; see package server for the UnifiedPush rules built on it). Each is a +// per-topic fact with a merge rule, so the nodes can write it in any order: the earliest first +// observation, the latest subscriber activity, and the last written rate visitor. +// +// With a store, a sync worker writes the topics changed here (a dirty set, so a round costs +// what changed, not the table) every syncInterval and reads back the rows changed since the last +// round, applying the same merge rules in memory; at startup it loads every row. A publish +// never reads the database (Get is memory only); Refresh is the one cold-path read, for a +// decision too final to make from a stale row. Rows not written for pruneAfter are pruned from +// the table and from memory alike. +// +// The worker is the service's own, not the usage counters' (package limit): a stuck topic sync +// must not stall the daily limits. Timestamps are the node's clock (unix seconds), so the nodes +// must keep their clocks in sync (NTP); a skew beyond readOverlap may skip rows in the read-back. +package topic + +import ( + "context" + "slices" + "sync" + "time" + + "heckel.io/ntfy/v2/log" +) + +const ( + tag = "topic" + + // syncInterval is how often a node writes its changed topics and reads the others'; a change + // reaches the other nodes within about two intervals (one to be written, one to be read) + syncInterval = 5 * time.Second + + // statementTimeout bounds each store call of a sync. A row's updated_at is the writing node's + // time before its statement, so it may commit up to this much later; reads re-read readOverlap + // behind the watermark so such a row (or one stamped by a node clock a little behind ours) is + // not skipped (the merge rules make a re-read harmless). + statementTimeout = 5 * time.Second + readOverlap = statementTimeout + time.Second + + // loadWait is how long Start waits for the startup load before serving from memory, while + // the load goes on in the background; loadTimeout bounds one load attempt + loadWait = 30 * time.Second + loadTimeout = 2 * time.Minute + + // pruneInterval is how often a node deletes rows not written for pruneAfter, from the table + // and from memory. pruneAfter must be well above the server's 404 threshold for an abandoned + // UnifiedPush topic (7 days), so the row that starts the countdown is still there when the + // threshold is reached. + pruneInterval = time.Hour + pruneAfter = 30 * 24 * time.Hour + + // keepaliveInterval is how often an open stream's subscriber activity is recorded at most + // (see Keepalive), so a long-lived stream costs one write per interval, not one per call + keepaliveInterval = 5 * time.Minute + + // maxNoted caps the topics known only from refused pushes (see Note), so a flood of pushes to + // topics nobody subscribed to cannot fill the memory of every node, or the table. Such a topic + // is only an entry of ~100 bytes, but it is kept for pruneAfter. + maxNoted = 100_000 +) + +// Service keeps the topics' state in memory and syncs it through the store; see the package comment +type Service struct { + conf *Config + now func() time.Time // Injectable clock, for tests + store *Store // Nil: memory only + entries map[string]*entry // By topic + dirty map[string]struct{} // Topics with an unwritten change + noted int // Entries without subscriber activity so far, capped at maxNoted + watermark int64 // Time of the last read, unix seconds; the sync worker's only + prunedAt int64 // Time of the last prune, unix seconds; the sync worker's only + interval time.Duration // How often the worker syncs (syncInterval); tests change it + ctx context.Context // Done on Close; the sync worker's store calls derive from it, so Close never waits on a stuck one + cancel context.CancelFunc + closeOnce sync.Once + wg sync.WaitGroup + mu sync.Mutex // Protects entries, dirty and noted +} + +// New creates the service +func New(conf *Config) *Service { + return newService(conf, time.Now) +} + +func newService(conf *Config, now func() time.Time) *Service { + s := &Service{ + conf: conf, + now: now, + store: conf.Store, + entries: make(map[string]*entry), + dirty: make(map[string]struct{}), + interval: syncInterval, + } + s.ctx, s.cancel = context.WithCancel(context.Background()) + return s +} + +// Start loads the current rows from the store, if there is one, and starts the sync worker. It +// returns once the load is done, or after loadWait (or when ctx is done), whichever is first; +// the load then goes on in the background while the service serves from memory. +func (s *Service) Start(ctx context.Context) error { + loaded := make(chan struct{}) + s.wg.Add(1) + go s.runSync(loaded) + if s.store == nil { + return nil + } + select { + case <-loaded: + case <-time.After(loadWait): + log.Tag(tag).Warn("Loading topics takes longer than %v; serving from memory, loading in the background", loadWait) + case <-ctx.Done(): + log.Tag(tag).Warn("Loading topics was not done in time; serving from memory, loading in the background") + } + return nil +} + +// Close stops the sync worker, cancelling its store call (a stuck load or sync), and writes what +// is unwritten. The store's database is its owner's to close (the user manager's, or the shared +// pool), after this. +func (s *Service) Close() { + s.closeOnce.Do(func() { + s.cancel() + s.wg.Wait() + if s.store == nil { + return + } + ctx, cancel := context.WithTimeout(context.Background(), statementTimeout) + defer cancel() + s.write(ctx) + }) +} + +// SetRateVisitor makes rv the topic's rate visitor and records subscriber activity now. The last +// one set, on any node, wins. +func (s *Service) SetRateVisitor(topic string, rv *RateVisitor) { + now := s.now().Unix() + s.mu.Lock() + defer s.mu.Unlock() + e := s.touchNoLock(topic, now) + e.userID, e.ip = rv.UserID, rv.IP + e.write = writeRateVisitor +} + +// Touch records subscriber activity on the topic now (a subscribe or poll by someone who may not +// pay, or a stream that closed), keeping its rate visitor +func (s *Service) Touch(topic string) { + now := s.now().Unix() + s.mu.Lock() + defer s.mu.Unlock() + s.touchNoLock(topic, now) +} + +// Keepalive records that a stream is still open on the topic, as a Touch at most once per +// keepaliveInterval, so the other nodes see the subscriber as active up to that much late +func (s *Service) Keepalive(topic string) { + now := s.now().Unix() + s.mu.Lock() + defer s.mu.Unlock() + if e, ok := s.entries[topic]; ok && now-e.subscribedAt < int64(keepaliveInterval/time.Second) { + return + } + s.touchNoLock(topic, now) +} + +// Note makes a topic this node does not know yet known, as first seen now, e.g. on a refused +// push, so that an abandoned topic can be told apart from a new one. It is insert-only: a known +// topic is left alone, so a push never counts as subscriber activity and never revives a +// rate visitor. At most maxNoted such topics are kept. +func (s *Service) Note(topic string) { + now := s.now().Unix() + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.entries[topic]; ok || s.noted >= maxNoted { + return + } + s.entries[topic] = &entry{ + firstSeenAt: now, + updatedAt: now, + write: writeTouch, + } + s.dirty[topic] = struct{}{} + s.noted++ +} + +// Get returns the topic's state as this node knows it, from memory, or nil if it is unknown +func (s *Service) Get(topic string) *Topic { + s.mu.Lock() + defer s.mu.Unlock() + return s.getNoLock(topic) +} + +// Refresh reads the topic's row from the store (a primary key read, with the caller's context), +// merges it into memory, and returns the topic's state like Get. It is for a decision too final +// to make from a stale row, on a cold path only. Without a store it is Get. +func (s *Service) Refresh(ctx context.Context, topic string) (*Topic, error) { + if s.store == nil { + return s.Get(topic), nil + } + r, err := s.store.readOne(ctx, topic) + if err != nil { + return nil, err + } + s.mu.Lock() + defer s.mu.Unlock() + if r != nil { + s.applyNoLock(r) + } + return s.getNoLock(topic), nil +} + +// RateVisitorUserIDs returns the distinct user IDs of the rate visitors of the topics with +// subscriber activity since the given time, sorted, e.g. to refresh their users' tiers +func (s *Service) RateVisitorUserIDs(since time.Time) []string { + cutoff := since.Unix() + ids := make(map[string]struct{}) + s.mu.Lock() + for _, e := range s.entries { + if e.userID != "" && e.subscribedAt >= cutoff { + ids[e.userID] = struct{}{} + } + } + s.mu.Unlock() + sorted := make([]string, 0, len(ids)) + for id := range ids { + sorted = append(sorted, id) + } + slices.Sort(sorted) + return sorted +} + +// Size returns the number of topics the service knows +func (s *Service) Size() int { + s.mu.Lock() + defer s.mu.Unlock() + return len(s.entries) +} + +// SetForTest replaces the topic's state with the given one, as if it had been observed here, and +// has it written. It exists so tests can age a topic; it is not used in production code. +func (s *Service) SetForTest(topic string, t *Topic) { + now := s.now().Unix() + s.mu.Lock() + defer s.mu.Unlock() + e := s.entryNoLock(topic, now) + e.firstSeenAt, e.subscribedAt = t.FirstSeenAt.Unix(), 0 + if !t.SubscribedAt.IsZero() { + e.subscribedAt = t.SubscribedAt.Unix() + } + e.userID, e.ip = "", "" + if t.RateVisitor != nil { + e.userID, e.ip = t.RateVisitor.UserID, t.RateVisitor.IP + } + e.updatedAt = now + e.write = max(e.write, writeTouch) + s.dirty[topic] = struct{}{} + s.countNotedNoLock() +} + +// runSync loads the current rows (if there is a store), retrying until it succeeds, then syncs +// every interval until Close. It is the only caller of the store's writes, so writes and reads +// never interleave. +func (s *Service) runSync(loaded chan struct{}) { + defer s.wg.Done() + for s.store != nil && !s.load() { + select { + case <-s.ctx.Done(): + return + case <-time.After(s.interval): + } + } + close(loaded) + ticker := time.NewTicker(s.interval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + s.sync() + case <-s.ctx.Done(): + return + } + } +} + +// load reads every row not pruned yet, and reports whether it succeeded +func (s *Service) load() bool { + ctx, cancel := context.WithTimeout(s.ctx, loadTimeout) + defer cancel() + started := time.Now() + now := s.now().Unix() + n, err := s.read(ctx, now-int64(pruneAfter/time.Second)) + if err != nil { + log.Tag(tag).Err(err).Warn("Cannot load topics; retrying in %v", s.interval) + return false + } + s.watermark, s.prunedAt = now, now + log.Tag(tag).Info("Loaded %d topic(s) in %v", n, time.Since(started).Round(time.Millisecond)) + return true +} + +// sync is one round of the sync worker: write the topics changed here, read back the rows changed +// since the last round, and prune once in a while, in memory and (with a store) in the table. +// Every error is logged and leaves the state as it was, for the next round. +func (s *Service) sync() { + now := s.now().Unix() + if s.store != nil { + ctx, cancel := context.WithTimeout(s.ctx, statementTimeout) + s.write(ctx) + cancel() + ctx, cancel = context.WithTimeout(s.ctx, statementTimeout) + _, err := s.read(ctx, s.watermark-int64(readOverlap/time.Second)) + cancel() + if err != nil { + log.Tag(tag).Err(err).Warn("Cannot read topics; those of other nodes are stale until the next round") + } else { + s.watermark = now + } + } + if now-s.prunedAt < int64(pruneInterval/time.Second) { + return + } + s.prunedAt = now + before := now - int64(pruneAfter/time.Second) + s.prune(before) + if s.store == nil { + return + } + ctx, cancel := context.WithTimeout(s.ctx, statementTimeout) + deleted, err := s.store.prune(ctx, before) + cancel() + if err != nil { + log.Tag(tag).Err(err).Warn("Cannot prune topics") + } else { + log.Tag(tag).Debug("Pruned %d topic row(s)", deleted) + } +} + +// write writes the topics changed here, as full rows. A failed write marks them changed again, +// for the next round; rows are state, not deltas, so rewriting the ones that were written is +// harmless. +func (s *Service) write(ctx context.Context) { + rows := s.dirtyRows() + if len(rows) == 0 { + return + } + now := s.now().Unix() + if err := s.store.write(ctx, rows, now); err != nil { + s.markDirty(rows) + log.Tag(tag).Err(err).Warn("Cannot write %d topic(s); retrying in the next round", len(rows)) + return + } + s.markWritten(rows, now) +} + +// dirtyRows returns the rows of the topics changed here, and marks them written +func (s *Service) dirtyRows() []row { + s.mu.Lock() + defer s.mu.Unlock() + rows := make([]row, 0, len(s.dirty)) + for topic := range s.dirty { + e, ok := s.entries[topic] + if !ok { + continue + } + rows = append(rows, row{ + topic: topic, + firstSeenAt: e.firstSeenAt, + subscribedAt: e.subscribedAt, + userID: e.userID, + ip: e.ip, + }) + e.write = writeNone + } + clear(s.dirty) + return rows +} + +// markDirty marks the topics of rows that failed to write as changed again. A row with a rate +// visitor is marked as if it were set here, so the read-back keeps it until it is written. +func (s *Service) markDirty(rows []row) { + s.mu.Lock() + defer s.mu.Unlock() + for _, r := range rows { + e, ok := s.entries[r.topic] + if !ok { + continue + } + write := writeTouch + if r.ip != "" { + write = writeRateVisitor + } + e.write = max(e.write, write) + s.dirty[r.topic] = struct{}{} + } +} + +// markWritten stamps the written rows' topics with the write time, which the prune goes by +func (s *Service) markWritten(rows []row, now int64) { + s.mu.Lock() + defer s.mu.Unlock() + for _, r := range rows { + if e, ok := s.entries[r.topic]; ok && e.updatedAt < now { + e.updatedAt = now + } + } +} + +// read applies the rows written since the given time to memory, and returns how many it applied +func (s *Service) read(ctx context.Context, since int64) (int, error) { + applied := 0 + err := s.store.readSince(ctx, since, func(r *row) { + s.mu.Lock() + defer s.mu.Unlock() + s.applyNoLock(r) + applied++ + }) + return applied, err +} + +// applyNoLock merges a row read from the store into memory, by the store's own rules: the earliest +// first observation, the latest subscriber activity, the row's rate visitor unless it has none or +// one was set here and not written yet (that one is newer). The caller must hold s.mu. +func (s *Service) applyNoLock(r *row) { + e, ok := s.entries[r.topic] + if !ok { + e = &entry{firstSeenAt: r.firstSeenAt} + s.entries[r.topic] = e + if r.subscribedAt == 0 { + s.noted++ + } + } else if r.firstSeenAt < e.firstSeenAt { + e.firstSeenAt = r.firstSeenAt + } + if r.subscribedAt > e.subscribedAt { + if e.subscribedAt == 0 && ok { + s.noted-- + } + e.subscribedAt = r.subscribedAt + } + if r.ip != "" && e.write != writeRateVisitor { + e.userID, e.ip = r.userID, r.ip + } + if r.updatedAt > e.updatedAt { + e.updatedAt = r.updatedAt + } +} + +// prune forgets the topics not written since the given time (unix seconds), unless they have an +// unwritten change (without a store, nothing is ever written, so there is nothing to wait for) +func (s *Service) prune(before int64) { + s.mu.Lock() + defer s.mu.Unlock() + for topic, e := range s.entries { + if (e.write == writeNone || s.store == nil) && e.updatedAt < before { + delete(s.entries, topic) + } + } + s.countNotedNoLock() +} + +// touchNoLock records subscriber activity on the topic at now, creating the entry if needed, and +// returns the entry, marked to be written. The caller must hold s.mu. +func (s *Service) touchNoLock(topic string, now int64) *entry { + e, ok := s.entries[topic] + if !ok { + e = &entry{firstSeenAt: now} + s.entries[topic] = e + } else if e.subscribedAt == 0 { + s.noted-- // Known from pushes only until now + } + e.subscribedAt = max(e.subscribedAt, now) + e.updatedAt = max(e.updatedAt, now) + e.write = max(e.write, writeTouch) + s.dirty[topic] = struct{}{} + return e +} + +// entryNoLock returns the topic's entry, creating it as first seen at now if needed. The caller +// must hold s.mu. +func (s *Service) entryNoLock(topic string, now int64) *entry { + e, ok := s.entries[topic] + if !ok { + e = &entry{firstSeenAt: now} + s.entries[topic] = e + } + return e +} + +// getNoLock returns the topic's state, or nil if it is unknown. The caller must hold s.mu. +func (s *Service) getNoLock(topic string) *Topic { + e, ok := s.entries[topic] + if !ok { + return nil + } + t := &Topic{ + FirstSeenAt: time.Unix(e.firstSeenAt, 0).UTC(), + } + if e.subscribedAt != 0 { + t.SubscribedAt = time.Unix(e.subscribedAt, 0).UTC() + } + if e.ip != "" { + t.RateVisitor = &RateVisitor{ + UserID: e.userID, + IP: e.ip, + } + } + return t +} + +// countNotedNoLock recounts the topics known from pushes only, after a change that may have +// made or unmade several at once (a prune, SetForTest). The caller must hold s.mu. +func (s *Service) countNotedNoLock() { + s.noted = 0 + for _, e := range s.entries { + if e.subscribedAt == 0 { + s.noted++ + } + } +} diff --git a/topic/service_test.go b/topic/service_test.go new file mode 100644 index 00000000..c7e824db --- /dev/null +++ b/topic/service_test.go @@ -0,0 +1,423 @@ +package topic + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// testClock is a clock the tests move by hand +type testClock struct { + t time.Time + mu sync.Mutex +} + +func (c *testClock) now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.t +} + +func (c *testClock) advance(d time.Duration) { + c.mu.Lock() + defer c.mu.Unlock() + c.t = c.t.Add(d) +} + +func newTestClock() *testClock { + return &testClock{t: time.Date(2026, 10, 25, 12, 0, 0, 0, time.UTC)} +} + +func newTestServiceWithClock(t *testing.T) (*Service, *testClock) { + t.Helper() + clock := newTestClock() + return newService(&Config{}, clock.now), clock +} + +// startTestService starts a service on the store with the given clock, whose worker never ticks +// on its own: the tests sync by hand. It is closed when the test ends. +func startTestService(t *testing.T, store *Store, now func() time.Time) *Service { + t.Helper() + s := newService(&Config{Store: store}, now) + s.interval = time.Hour + require.Nil(t, s.Start(context.Background())) + t.Cleanup(s.Close) + return s +} + +func TestService_SetTouchNote(t *testing.T) { + s, clock := newTestServiceWithClock(t) + require.Nil(t, s.Get("upA")) + + // A refused push makes the topic known: first seen now, never subscribed, no rate visitor; a + // second one changes nothing + s.Note("upA") + noted := clock.now() + clock.advance(time.Hour) + s.Note("upA") + require.Equal(t, &Topic{FirstSeenAt: noted}, s.Get("upA")) + + // A subscriber becomes the rate visitor, and counts as subscriber activity; the latest one wins + s.SetRateVisitor("upA", &RateVisitor{IP: "1.2.3.4"}) + s.SetRateVisitor("upA", &RateVisitor{UserID: "u_1", IP: "1.2.3.4"}) + set := clock.now() + require.Equal(t, &Topic{FirstSeenAt: noted, SubscribedAt: set, RateVisitor: userRateVisitor()}, s.Get("upA")) + + // A touch keeps the rate visitor and moves the subscriber activity + clock.advance(time.Minute) + s.Touch("upA") + touched := clock.now() + require.Equal(t, &Topic{FirstSeenAt: noted, SubscribedAt: touched, RateVisitor: userRateVisitor()}, s.Get("upA")) + + // A refused push to a known topic changes nothing, however old its subscriber activity + clock.advance(48 * time.Hour) + s.Note("upA") + require.Equal(t, &Topic{FirstSeenAt: noted, SubscribedAt: touched, RateVisitor: userRateVisitor()}, s.Get("upA")) + + // A touch or a subscribe of an unknown topic makes it known, first seen and subscribed now + s.Touch("upB") + require.Equal(t, &Topic{FirstSeenAt: clock.now(), SubscribedAt: clock.now()}, s.Get("upB")) + s.SetRateVisitor("upC", &RateVisitor{IP: "1.2.3.4"}) + require.Equal(t, &Topic{FirstSeenAt: clock.now(), SubscribedAt: clock.now(), RateVisitor: ipRateVisitor()}, s.Get("upC")) + require.Equal(t, 3, s.Size()) +} + +func TestService_Keepalive_Throttled(t *testing.T) { + s, clock := newTestServiceWithClock(t) + s.SetRateVisitor("upA", ipRateVisitor()) + set := clock.now() + clock.advance(keepaliveInterval - time.Second) + s.Keepalive("upA") + require.Equal(t, set, s.Get("upA").SubscribedAt) + clock.advance(time.Second) + s.Keepalive("upA") + require.Equal(t, clock.now(), s.Get("upA").SubscribedAt) + require.Equal(t, ipRateVisitor(), s.Get("upA").RateVisitor) + + // A keepalive of an unknown topic is a touch + s.Keepalive("upB") + require.Equal(t, &Topic{FirstSeenAt: clock.now(), SubscribedAt: clock.now()}, s.Get("upB")) +} + +func TestService_Note_Bounded(t *testing.T) { + // Topics known only from refused pushes are capped; a subscriber frees a slot + s, _ := newTestServiceWithClock(t) + for i := 0; i < maxNoted; i++ { + s.Note(fmt.Sprintf("up%012d", i)) + } + require.Equal(t, maxNoted, s.Size()) + s.Note("upOneTooMany") + require.Nil(t, s.Get("upOneTooMany")) + s.Touch("up000000000000") + s.Note("upOneTooMany") + require.NotNil(t, s.Get("upOneTooMany")) + s.Note("upTwoTooMany") + require.Nil(t, s.Get("upTwoTooMany")) +} + +func TestService_Prune(t *testing.T) { + // Topics are pruned pruneAfter after their last write, never while they have an unwritten change + s, clock := newTestServiceWithClock(t) + s.Note("upA") + clock.advance(time.Hour) + s.Touch("upB") + clock.advance(pruneAfter - 30*time.Minute) + s.prune(clock.now().Unix() - int64(pruneAfter/time.Second)) + require.Nil(t, s.Get("upA")) + require.NotNil(t, s.Get("upB")) + + // A pruned topic may be noted again + for i := 0; i < maxNoted; i++ { + s.Note(fmt.Sprintf("up%012d", i)) + } + s.Note("upOneTooMany") + require.Nil(t, s.Get("upOneTooMany")) + clock.advance(pruneAfter + time.Hour) + s.prune(clock.now().Unix() - int64(pruneAfter/time.Second)) + s.Note("upOneTooMany") + require.NotNil(t, s.Get("upOneTooMany")) +} + +func TestService_RateVisitorUserIDs(t *testing.T) { + s, clock := newTestServiceWithClock(t) + s.SetRateVisitor("upA", &RateVisitor{UserID: "u_1", IP: "1.2.3.4"}) + s.SetRateVisitor("upB", &RateVisitor{IP: "1.2.3.4"}) + s.SetRateVisitor("upC", &RateVisitor{UserID: "u_2", IP: "5.6.7.8"}) + s.SetRateVisitor("upD", &RateVisitor{UserID: "u_1", IP: "5.6.7.8"}) + s.Note("upE") + clock.advance(25 * time.Hour) + s.Touch("upA") + require.Equal(t, []string{"u_1"}, s.RateVisitorUserIDs(clock.now().Add(-24*time.Hour))) + require.Equal(t, []string{"u_1", "u_2"}, s.RateVisitorUserIDs(clock.now().Add(-48*time.Hour))) +} + +func TestService_SetForTest(t *testing.T) { + s, clock := newTestServiceWithClock(t) + s.SetRateVisitor("upA", ipRateVisitor()) + aged := &Topic{ + FirstSeenAt: clock.now().Add(-48 * time.Hour), + SubscribedAt: clock.now().Add(-25 * time.Hour), + RateVisitor: ipRateVisitor(), + } + s.SetForTest("upA", aged) + require.Equal(t, aged, s.Get("upA")) + s.SetForTest("upB", &Topic{FirstSeenAt: clock.now().Add(-8 * 24 * time.Hour)}) + require.Equal(t, &Topic{FirstSeenAt: clock.now().Add(-8 * 24 * time.Hour)}, s.Get("upB")) +} + +func TestService_Persist_SharedAcrossNodes(t *testing.T) { + // A rate visitor set on one node is the rate visitor on the other after a sync; a touch on the + // other node moves the subscriber activity for both, and keeps the rate visitor + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + a := startTestService(t, newStore(), clock.now) + b := startTestService(t, newStore(), clock.now) + a.SetRateVisitor("upA", userRateVisitor()) + set := clock.now() + a.sync() + b.sync() + require.Equal(t, &Topic{FirstSeenAt: set, SubscribedAt: set, RateVisitor: userRateVisitor()}, b.Get("upA")) + clock.advance(time.Hour) + b.Touch("upA") + b.sync() + a.sync() + require.Equal(t, &Topic{FirstSeenAt: set, SubscribedAt: clock.now(), RateVisitor: userRateVisitor()}, a.Get("upA")) + require.Equal(t, a.Get("upA"), b.Get("upA")) + }) +} + +func TestService_Persist_SurvivesRestart(t *testing.T) { + // Close writes what is unwritten; a start loads every row + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + s := startTestService(t, newStore(), clock.now) + s.SetRateVisitor("upA", ipRateVisitor()) + s.Note("upB") + s.Close() + s = startTestService(t, newStore(), clock.now) + require.Equal(t, &Topic{FirstSeenAt: clock.now(), SubscribedAt: clock.now(), RateVisitor: ipRateVisitor()}, s.Get("upA")) + require.Equal(t, &Topic{FirstSeenAt: clock.now()}, s.Get("upB")) + require.Equal(t, 2, s.Size()) + }) +} + +func TestService_Persist_NoteDoesNotReviveRateVisitor(t *testing.T) { + // A refused push on a node that does not know the topic neither moves its subscriber activity + // nor touches its rate visitor on the other nodes: it is insert-only + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + a := startTestService(t, newStore(), clock.now) + b := startTestService(t, newStore(), clock.now) + a.SetRateVisitor("upA", userRateVisitor()) + set := clock.now() + a.sync() + clock.advance(25 * time.Hour) + b.Note("upA") // b has not read a's row yet + b.sync() + a.sync() + b.sync() + for _, s := range []*Service{a, b} { + require.Equal(t, &Topic{FirstSeenAt: set, SubscribedAt: set, RateVisitor: userRateVisitor()}, s.Get("upA")) + } + }) +} + +func TestService_Persist_TouchAfterPruneRestoresRow(t *testing.T) { + // A touch written after another node pruned the row recreates it with the rate visitor: every + // write carries the full row + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + store := newStore() + s := startTestService(t, store, clock.now) + s.SetRateVisitor("upA", userRateVisitor()) + set := clock.now() + s.sync() + deleted, err := store.prune(context.Background(), clock.now().Unix()+1) + require.Nil(t, err) + require.Equal(t, int64(1), deleted) + clock.advance(time.Hour) + s.Touch("upA") + s.sync() + require.Equal(t, row{topic: "upA", firstSeenAt: set.Unix(), subscribedAt: clock.now().Unix(), userID: "u_1", ip: "1.2.3.4", updatedAt: clock.now().Unix()}, readAll(t, store, 0)["upA"]) + }) +} + +func TestService_Persist_ReadBackKeepsUnwrittenRateVisitor(t *testing.T) { + // A rate visitor set here and not written yet is not replaced by a row read back; once + // written, it is the one every node has (last write wins) + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + a := startTestService(t, newStore(), clock.now) + b := startTestService(t, newStore(), clock.now) + a.SetRateVisitor("upA", ipRateVisitor()) + a.sync() + b.SetRateVisitor("upA", &RateVisitor{IP: "5.6.7.8"}) + _, err := b.read(context.Background(), 0) + require.Nil(t, err) + require.Equal(t, &RateVisitor{IP: "5.6.7.8"}, b.Get("upA").RateVisitor) + b.sync() + a.sync() + require.Equal(t, &RateVisitor{IP: "5.6.7.8"}, a.Get("upA").RateVisitor) + }) +} + +func TestService_Persist_MergesFirstSeenAndSubscribed(t *testing.T) { + // Two nodes that know a topic differently converge on the earliest first observation and the + // latest subscriber activity, whichever writes first + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + a := startTestService(t, newStore(), clock.now) + b := startTestService(t, newStore(), clock.now) + a.Note("upA") + noted := clock.now() + clock.advance(time.Hour) + b.Touch("upA") + touched := clock.now() + b.sync() + a.sync() + b.sync() + for _, s := range []*Service{a, b} { + require.Equal(t, &Topic{FirstSeenAt: noted, SubscribedAt: touched}, s.Get("upA")) + } + }) +} + +func TestService_Persist_Refresh(t *testing.T) { + // Refresh reads the row from the store, so a node that has not read back a fresher row yet + // decides on the fresh one, and keeps it + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + a := startTestService(t, newStore(), clock.now) + b := startTestService(t, newStore(), clock.now) + a.SetRateVisitor("upA", ipRateVisitor()) + set := clock.now() + a.sync() + require.Nil(t, b.Get("upA")) + topic, err := b.Refresh(context.Background(), "upA") + require.Nil(t, err) + require.Equal(t, &Topic{FirstSeenAt: set, SubscribedAt: set, RateVisitor: ipRateVisitor()}, topic) + require.Equal(t, topic, b.Get("upA")) + + // A fresher row replaces what b knows + clock.advance(time.Hour) + a.Touch("upA") + a.sync() + topic, err = b.Refresh(context.Background(), "upA") + require.Nil(t, err) + require.Equal(t, clock.now(), topic.SubscribedAt) + require.Equal(t, clock.now(), b.Get("upA").SubscribedAt) + + // An unknown topic stays unknown; a failed read is an error + topic, err = b.Refresh(context.Background(), "upB") + require.Nil(t, err) + require.Nil(t, topic) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = b.Refresh(ctx, "upA") + require.NotNil(t, err) + }) +} + +func TestService_Persist_WriteFailureRetries(t *testing.T) { + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + store := newStore() + s := startTestService(t, store, clock.now) + s.SetRateVisitor("upA", ipRateVisitor()) + writeChunk := store.writeChunk + store.writeChunk = func(ctx context.Context, rows []row, now int64) error { + return errors.New("database is down") + } + s.sync() + require.Empty(t, readAll(t, store, 0)) + store.writeChunk = writeChunk + s.sync() + require.Equal(t, "1.2.3.4", readAll(t, store, 0)["upA"].ip) + }) +} + +func TestService_Persist_PrunesTableAndMemory(t *testing.T) { + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + clock := newTestClock() + store := newStore() + s := startTestService(t, store, clock.now) + s.Note("upA") + s.sync() + clock.advance(pruneAfter + 2*time.Hour) + s.Note("upB") + s.sync() // Prunes: an hour has passed since the last one + require.Nil(t, s.Get("upA")) + require.NotNil(t, s.Get("upB")) + require.Equal(t, []string{"upB"}, topicsOf(readAll(t, store, 0))) + + // A topic with an unwritten change is never pruned from memory, however old its last write + s.SetRateVisitor("upC", ipRateVisitor()) + store.writeChunk = func(ctx context.Context, rows []row, now int64) error { + return errors.New("database is down") + } + s.sync() + s.mu.Lock() + s.entries["upC"].updatedAt = 1 + s.mu.Unlock() + clock.advance(pruneInterval) + s.sync() + require.NotNil(t, s.Get("upC")) + }) +} + +func TestService_WithoutStore(t *testing.T) { + s := New(&Config{}) + require.Nil(t, s.Start(context.Background())) + s.Note("upA") + topic, err := s.Refresh(context.Background(), "upA") + require.Nil(t, err) + require.Equal(t, s.Get("upA"), topic) + s.Close() +} + +func TestService_Close_CancelsStalledLoad(t *testing.T) { + // Close does not wait for a startup load stuck on the database + store := NewSQLiteStore(newTestUserDB(t)) + pool := store.db.Primary() + pool.SetMaxOpenConns(1) + conn, err := pool.Conn(context.Background()) // Holds the only connection: the load waits for it + require.Nil(t, err) + defer conn.Close() + s := newService(&Config{Store: store}, time.Now) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + require.Nil(t, s.Start(ctx)) + require.Eventually(t, func() bool { + return pool.Stats().WaitCount > 0 + }, time.Second, time.Millisecond) + closed := make(chan struct{}) + go func() { + s.Close() + close(closed) + }() + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("Close waits for the stalled load") + } +} + +// userRateVisitor returns the rate visitor of user u_1 at 1.2.3.4 +func userRateVisitor() *RateVisitor { + return &RateVisitor{ + UserID: "u_1", + IP: "1.2.3.4", + } +} + +// ipRateVisitor returns the rate visitor of the anonymous subscriber at 1.2.3.4 +func ipRateVisitor() *RateVisitor { + return &RateVisitor{ + IP: "1.2.3.4", + } +} diff --git a/topic/store.go b/topic/store.go new file mode 100644 index 00000000..87f49f98 --- /dev/null +++ b/topic/store.go @@ -0,0 +1,124 @@ +package topic + +import ( + "context" + "database/sql" + "slices" + "strings" + + "heckel.io/ntfy/v2/db" +) + +const ( + // writeChunkRows caps the rows of one write statement (Postgres) or transaction (SQLite), so + // a big sync is a few bounded statements rather than one huge one + writeChunkRows = 1000 + + // pruneChunkRows caps the rows of one prune statement, so no statement holds locks for long + pruneChunkRows = 10000 +) + +// row is one topic's row of the topic table. On the way to the store, an empty ip keeps the +// stored rate visitor (see the upsert queries); updated_at is stamped by the store. +type row struct { + topic string + firstSeenAt int64 // Unix seconds + subscribedAt int64 // Unix seconds; 0 if never + userID string + ip string + updatedAt int64 // Unix seconds; only on the way back +} + +// Store reads and writes the topic table, on SQLite (the user database) or Postgres. Rows are +// stamped (updated_at) with the writing node's clock, like the reader's watermark, so the +// database's clock does not matter. The service's sync worker calls it, and a cold-path +// Refresh. +type Store struct { + db *db.DB + queries *queries + writeChunk func(ctx context.Context, rows []row, now int64) error // Dialect specific: unnest on Postgres, a prepared statement per row on SQLite +} + +// queries are the dialect's SQL statements +type queries struct { + selectSince string + selectOne string + upsert string + prune string +} + +// write merges the rows into the table (see the upsert queries: first_seen_at keeps the earliest, +// subscribed_at the latest, the rate visitor the last written one, unless the row has none), +// sorted by topic in chunks so two nodes lock overlapping topics in the same order, stamped with +// now (unix seconds). Rows are state, not deltas, so a failed chunk is simply written again. +func (s *Store) write(ctx context.Context, rows []row, now int64) error { + slices.SortFunc(rows, func(a, b row) int { + return strings.Compare(a.topic, b.topic) + }) + for chunk := range slices.Chunk(rows, writeChunkRows) { + if err := s.writeChunk(ctx, chunk, now); err != nil { + return err + } + } + return nil +} + +// readSince calls fn for every row written at or after the given time (unix seconds), streaming +// them from the primary (a replica may not have this node's own writes yet) +func (s *Store) readSince(ctx context.Context, since int64, fn func(r *row)) error { + rows, err := s.db.QueryContext(ctx, s.queries.selectSince, since) + if err != nil { + return err + } + defer rows.Close() + var r row + for rows.Next() { + if err := scanRow(rows, &r); err != nil { + return err + } + fn(&r) + } + return rows.Err() +} + +// readOne returns the topic's row from the primary, or nil if there is none +func (s *Store) readOne(ctx context.Context, topic string) (*row, error) { + rows, err := s.db.QueryContext(ctx, s.queries.selectOne, topic) + if err != nil { + return nil, err + } + defer rows.Close() + if !rows.Next() { + return nil, rows.Err() + } + var r row + if err := scanRow(rows, &r); err != nil { + return nil, err + } + return &r, nil +} + +// prune deletes the rows not written since the given time (unix seconds), in chunks, and returns +// how many it deleted +func (s *Store) prune(ctx context.Context, before int64) (int64, error) { + var deleted int64 + for { + result, err := s.db.ExecContext(ctx, s.queries.prune, before, pruneChunkRows) + if err != nil { + return deleted, err + } + n, err := result.RowsAffected() + if err != nil { + return deleted, err + } + deleted += n + if n < pruneChunkRows { + return deleted, nil + } + } +} + +// scanRow reads the current row of a select into r; the selects list the columns in this order +func scanRow(rows *sql.Rows, r *row) error { + return rows.Scan(&r.topic, &r.firstSeenAt, &r.subscribedAt, &r.userID, &r.ip, &r.updatedAt) +} diff --git a/topic/store_postgres.go b/topic/store_postgres.go new file mode 100644 index 00000000..7520a215 --- /dev/null +++ b/topic/store_postgres.go @@ -0,0 +1,54 @@ +package topic + +import ( + "context" + + "heckel.io/ntfy/v2/db" +) + +// The upsert is one fixed statement per chunk (unnest), because SQL built per call bloats pgx's +// statement cache. Merge rules: the earliest first_seen_at, the latest subscribed_at, and the +// last written rate visitor, except that a row without one (no ip) keeps the stored one. +const ( + postgresSelectSinceQuery = `SELECT topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at FROM topic WHERE updated_at >= $1` + postgresSelectOneQuery = `SELECT topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at FROM topic WHERE topic = $1` + postgresUpsertQuery = ` + INSERT INTO topic AS t (topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at) + SELECT topic, f, s, u, i, $6::BIGINT + FROM unnest($1::TEXT[], $2::BIGINT[], $3::BIGINT[], $4::TEXT[], $5::TEXT[]) AS d(topic, f, s, u, i) + ON CONFLICT (topic) DO UPDATE SET + first_seen_at = LEAST(t.first_seen_at, EXCLUDED.first_seen_at), + subscribed_at = GREATEST(t.subscribed_at, EXCLUDED.subscribed_at), + rate_visitor_user_id = CASE WHEN EXCLUDED.rate_visitor_ip = '' THEN t.rate_visitor_user_id ELSE EXCLUDED.rate_visitor_user_id END, + rate_visitor_ip = CASE WHEN EXCLUDED.rate_visitor_ip = '' THEN t.rate_visitor_ip ELSE EXCLUDED.rate_visitor_ip END, + updated_at = EXCLUDED.updated_at + ` + postgresPruneQuery = `DELETE FROM topic WHERE updated_at < $1 AND topic IN (SELECT topic FROM topic WHERE updated_at < $1 LIMIT $2)` // The outer age check is rechecked on a row updated while the delete waits for it +) + +// NewPostgresStore returns the store on the shared pool; the topic table is the user schema's +// (see package user, which creates and migrates it) +func NewPostgresStore(d *db.DB) *Store { + s := &Store{ + db: d, + queries: &queries{ + selectSince: postgresSelectSinceQuery, + selectOne: postgresSelectOneQuery, + upsert: postgresUpsertQuery, + prune: postgresPruneQuery, + }, + } + s.writeChunk = s.writePostgres + return s +} + +// writePostgres writes the rows in one statement, as arrays +func (s *Store) writePostgres(ctx context.Context, rows []row, now int64) error { + topics, userIDs, ips := make([]string, len(rows)), make([]string, len(rows)), make([]string, len(rows)) + firstSeenAts, subscribedAts := make([]int64, len(rows)), make([]int64, len(rows)) + for i, r := range rows { + topics[i], firstSeenAts[i], subscribedAts[i], userIDs[i], ips[i] = r.topic, r.firstSeenAt, r.subscribedAt, r.userID, r.ip + } + _, err := s.db.ExecContext(ctx, s.queries.upsert, topics, firstSeenAts, subscribedAts, userIDs, ips, now) + return err +} diff --git a/topic/store_sqlite.go b/topic/store_sqlite.go new file mode 100644 index 00000000..bfbc0cb0 --- /dev/null +++ b/topic/store_sqlite.go @@ -0,0 +1,59 @@ +package topic + +import ( + "context" + "database/sql" + + "heckel.io/ntfy/v2/db" +) + +// Same as the Postgres queries (see there), one row per statement: SQLite is a local file, so +// a prepared statement per row in one transaction costs no round trips +const ( + sqliteSelectSinceQuery = `SELECT topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at FROM topic WHERE updated_at >= ?` + sqliteSelectOneQuery = `SELECT topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at FROM topic WHERE topic = ?` + sqliteUpsertQuery = ` + INSERT INTO topic (topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT (topic) DO UPDATE SET + first_seen_at = MIN(first_seen_at, excluded.first_seen_at), + subscribed_at = MAX(subscribed_at, excluded.subscribed_at), + rate_visitor_user_id = CASE WHEN excluded.rate_visitor_ip = '' THEN rate_visitor_user_id ELSE excluded.rate_visitor_user_id END, + rate_visitor_ip = CASE WHEN excluded.rate_visitor_ip = '' THEN rate_visitor_ip ELSE excluded.rate_visitor_ip END, + updated_at = excluded.updated_at + ` + sqlitePruneQuery = `DELETE FROM topic WHERE updated_at < ?1 AND topic IN (SELECT topic FROM topic WHERE updated_at < ?1 LIMIT ?2)` +) + +// NewSQLiteStore returns the store on the user database (auth-file), which holds the topic +// table (see package user, which creates and migrates it) +func NewSQLiteStore(d *db.DB) *Store { + s := &Store{ + db: d, + queries: &queries{ + selectSince: sqliteSelectSinceQuery, + selectOne: sqliteSelectOneQuery, + upsert: sqliteUpsertQuery, + prune: sqlitePruneQuery, + }, + } + s.writeChunk = s.writeSQLite + return s +} + +// writeSQLite writes the rows in one transaction, one prepared statement per row +func (s *Store) writeSQLite(ctx context.Context, rows []row, now int64) error { + return db.ExecTx(s.db, func(tx *sql.Tx) error { + stmt, err := tx.PrepareContext(ctx, s.queries.upsert) + if err != nil { + return err + } + defer stmt.Close() + for _, r := range rows { + if _, err := stmt.ExecContext(ctx, r.topic, r.firstSeenAt, r.subscribedAt, r.userID, r.ip, now); err != nil { + return err + } + } + return nil + }) +} diff --git a/topic/store_test.go b/topic/store_test.go new file mode 100644 index 00000000..5f294eff --- /dev/null +++ b/topic/store_test.go @@ -0,0 +1,149 @@ +package topic + +import ( + "context" + "fmt" + "path/filepath" + "slices" + "testing" + + "github.com/stretchr/testify/require" + "heckel.io/ntfy/v2/db" + dbtest "heckel.io/ntfy/v2/db/test" + "heckel.io/ntfy/v2/user" +) + +// forEachBackend runs f against a fresh topic table on SQLite and on Postgres (if +// NTFY_TEST_DATABASE_URL is set). newStore returns a store on the same table every time, like a +// second node or a restarted one. +func forEachBackend(t *testing.T, f func(t *testing.T, newStore func() *Store)) { + t.Run("sqlite", func(t *testing.T) { + d := newTestUserDB(t) + f(t, func() *Store { + return NewSQLiteStore(d) + }) + }) + t.Run("postgres", func(t *testing.T) { + d := dbtest.CreateTestPostgres(t) + _, err := user.NewPostgresManager(d, &user.Config{}) // Creates the topic table + require.Nil(t, err) + f(t, func() *Store { + return NewPostgresStore(d) + }) + }) +} + +// newTestUserDB returns a fresh SQLite user database, which holds the topic table +func newTestUserDB(t testing.TB) *db.DB { + d, err := user.OpenSQLite(filepath.Join(t.TempDir(), "user.db")) + require.Nil(t, err) + manager, err := user.NewSQLiteManager(d, "", &user.Config{}) + require.Nil(t, err) + t.Cleanup(func() { + manager.Close() + }) + return d +} + +// readAll returns the rows written at or after since, by topic +func readAll(t *testing.T, s *Store, since int64) map[string]row { + t.Helper() + rows := make(map[string]row) + require.Nil(t, s.readSince(context.Background(), since, func(r *row) { + rows[r.topic] = *r + })) + return rows +} + +// topicsOf returns the map's topics, sorted +func topicsOf(rows map[string]row) []string { + topics := make([]string, 0, len(rows)) + for topic := range rows { + topics = append(topics, topic) + } + slices.Sort(topics) + return topics +} + +func TestStore_Write_MergeRules(t *testing.T) { + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + s := newStore() + write := func(r row, now int64) { + require.Nil(t, s.write(context.Background(), []row{r}, now)) + } + current := func() row { + return readAll(t, s, 0)["up1"] + } + + // A row without a rate visitor (a refused push) inserts as it is + write(row{topic: "up1", firstSeenAt: 100}, 101) + require.Equal(t, row{topic: "up1", firstSeenAt: 100, updatedAt: 101}, current()) + + // A subscribe sets the rate visitor and the subscriber activity; first_seen_at keeps the earliest + write(row{topic: "up1", firstSeenAt: 200, subscribedAt: 200, userID: "u_1", ip: "1.2.3.4"}, 201) + require.Equal(t, row{topic: "up1", firstSeenAt: 100, subscribedAt: 200, userID: "u_1", ip: "1.2.3.4", updatedAt: 201}, current()) + + // A row without a rate visitor keeps the stored one, and an older subscribed_at does not move it back + write(row{topic: "up1", firstSeenAt: 300, subscribedAt: 150}, 301) + require.Equal(t, row{topic: "up1", firstSeenAt: 100, subscribedAt: 200, userID: "u_1", ip: "1.2.3.4", updatedAt: 301}, current()) + + // A later subscribed_at moves it forward + write(row{topic: "up1", firstSeenAt: 300, subscribedAt: 250}, 302) + require.Equal(t, row{topic: "up1", firstSeenAt: 100, subscribedAt: 250, userID: "u_1", ip: "1.2.3.4", updatedAt: 302}, current()) + + // The last written rate visitor wins, whatever its subscribed_at; an earlier first_seen_at moves it back + write(row{topic: "up1", firstSeenAt: 50, subscribedAt: 240, ip: "5.6.7.8"}, 303) + require.Equal(t, row{topic: "up1", firstSeenAt: 50, subscribedAt: 250, ip: "5.6.7.8", updatedAt: 303}, current()) + }) +} + +func TestStore_Write_Chunked(t *testing.T) { + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + s := newStore() + rows := make([]row, 0) + for i := 0; i < writeChunkRows+5; i++ { + rows = append(rows, row{topic: fmt.Sprintf("up%012d", i), firstSeenAt: 100, subscribedAt: 100, ip: "1.2.3.4"}) + } + require.Nil(t, s.write(context.Background(), rows, 100)) + require.Len(t, readAll(t, s, 0), writeChunkRows+5) + }) +} + +func TestStore_ReadSince(t *testing.T) { + // The read-back is by write time, not by the topic's own times + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + s := newStore() + for i := 0; i < 5; i++ { + require.Nil(t, s.write(context.Background(), []row{{topic: fmt.Sprintf("up%d", i), firstSeenAt: 1000, subscribedAt: 1000}}, int64(100+i))) + } + require.Equal(t, []string{"up0", "up1", "up2", "up3", "up4"}, topicsOf(readAll(t, s, 0))) + require.Equal(t, []string{"up3", "up4"}, topicsOf(readAll(t, s, 103))) + require.Empty(t, readAll(t, s, 105)) + }) +} + +func TestStore_ReadOne(t *testing.T) { + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + s := newStore() + require.Nil(t, s.write(context.Background(), []row{{topic: "up1", firstSeenAt: 100, subscribedAt: 200, userID: "u_1", ip: "1.2.3.4"}}, 201)) + r, err := s.readOne(context.Background(), "up1") + require.Nil(t, err) + require.Equal(t, &row{topic: "up1", firstSeenAt: 100, subscribedAt: 200, userID: "u_1", ip: "1.2.3.4", updatedAt: 201}, r) + r, err = s.readOne(context.Background(), "up2") + require.Nil(t, err) + require.Nil(t, r) + }) +} + +func TestStore_Prune(t *testing.T) { + forEachBackend(t, func(t *testing.T, newStore func() *Store) { + s := newStore() + for i := 0; i < 5; i++ { + require.Nil(t, s.write(context.Background(), []row{{topic: fmt.Sprintf("up%d", i), firstSeenAt: 1}}, int64(100+i))) + } + deleted, err := s.prune(context.Background(), 103) + require.Nil(t, err) + require.Equal(t, int64(3), deleted) + require.Equal(t, []string{"up3", "up4"}, topicsOf(readAll(t, s, 0))) + }) +} diff --git a/topic/types.go b/topic/types.go new file mode 100644 index 00000000..57b3efd8 --- /dev/null +++ b/topic/types.go @@ -0,0 +1,45 @@ +package topic + +import ( + "time" +) + +// Config configures the service +type Config struct { + Store *Store // Where the rows are shared with the other nodes and kept across restarts; nil for memory only +} + +// RateVisitor is the subscriber who pays for the pushes to a UnifiedPush topic: its user, if it +// was logged in, and its IP, masked to the server's prefix bits. Which of the two is charged is +// decided on every push (see package server), so a tier change applies without a new subscribe. +type RateVisitor struct { + UserID string // Empty if the subscriber was anonymous + IP string // The masked IP, never empty +} + +// Topic is a topic's cluster-wide state as this node knows it (see Service.Get) +type Topic struct { + FirstSeenAt time.Time // Earliest observation on any node: a refused push, or the first subscribe + SubscribedAt time.Time // Last subscriber activity on any node; zero if there was none yet (the topic was only pushed to) + RateVisitor *RateVisitor // Who pays for the pushes; nil if nobody who may pay subscribed yet +} + +// write is what an entry still has to write to the store; a rate visitor set here outranks a touch +type write uint8 + +const ( + writeNone write = iota + writeTouch // The row as it is here (Touch, Keepalive, Note) + writeRateVisitor // A rate visitor set here (SetRateVisitor): until written, the read-back does not replace it +) + +// entry is one topic's row as this node knows it: the store's row as last read, merged with what +// happened here since (see Service.applyNoLock for the rules) +type entry struct { + firstSeenAt int64 // Unix seconds + subscribedAt int64 // Unix seconds; 0 if never + userID string // The rate visitor's user ID, empty if anonymous or none + ip string // The rate visitor's masked IP, empty if none + updatedAt int64 // When the row was last written, by any node (unix seconds); the prune goes by it + write write // Unwritten change, if any +} diff --git a/user/manager_postgres_schema.go b/user/manager_postgres_schema.go index d4944267..ace7fa05 100644 --- a/user/manager_postgres_schema.go +++ b/user/manager_postgres_schema.go @@ -146,15 +146,17 @@ const ( ` ) -// The topic table holds who pays for the pushes to a UnifiedPush topic and when it last had -// a subscriber, one row per topic (see package limit, which reads and writes it) +// The topic table is a topic's cluster-wide state: when it was first seen on any node, when a +// subscriber was last active, and who pays for the pushes to it (see package topic, which reads +// and writes it) const ( postgresCreateTopicTableQuery = ` CREATE TABLE IF NOT EXISTS topic ( topic TEXT NOT NULL PRIMARY KEY, - rate_visitor_key TEXT NOT NULL, + first_seen_at BIGINT NOT NULL, + subscribed_at BIGINT NOT NULL, rate_visitor_user_id TEXT NOT NULL, - seen_at BIGINT NOT NULL, + rate_visitor_ip TEXT NOT NULL, updated_at BIGINT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_topic_updated_at ON topic (updated_at); diff --git a/user/manager_sqlite_schema.go b/user/manager_sqlite_schema.go index f56e70cc..6d8dfdb8 100644 --- a/user/manager_sqlite_schema.go +++ b/user/manager_sqlite_schema.go @@ -387,15 +387,17 @@ const ( ` ) -// The topic table holds who pays for the pushes to a UnifiedPush topic and when it last had -// a subscriber, one row per topic (see package limit, which reads and writes it) +// The topic table is a topic's cluster-wide state: when it was first seen on any node, when a +// subscriber was last active, and who pays for the pushes to it (see package topic, which reads +// and writes it) const ( sqliteCreateTopicTableQuery = ` CREATE TABLE IF NOT EXISTS topic ( topic TEXT NOT NULL PRIMARY KEY, - rate_visitor_key TEXT NOT NULL, + first_seen_at INT NOT NULL, + subscribed_at INT NOT NULL, rate_visitor_user_id TEXT NOT NULL, - seen_at INT NOT NULL, + rate_visitor_ip TEXT NOT NULL, updated_at INT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_topic_updated_at ON topic (updated_at); diff --git a/user/manager_test.go b/user/manager_test.go index eeac14ca..63dda53c 100644 --- a/user/manager_test.go +++ b/user/manager_test.go @@ -2017,6 +2017,58 @@ func TestStoreUserByID(t *testing.T) { }) } +func TestStoreUsersByIDs(t *testing.T) { + // Users are looked up in bulk, with their tiers; unknown IDs are skipped, deleted users are + // returned as deleted, and more IDs than fit one statement are looked up in chunks + forEachBackend(t, func(t *testing.T, newManager newManagerFunc) { + manager := newTestManager(t, newManager, PermissionDenyAll) + require.Nil(t, manager.AddTier(&Tier{ + Code: "pro", + MessageLimit: 123, + })) + require.Nil(t, manager.AddUser("phil", "mypass", RoleAdmin, false)) + require.Nil(t, manager.AddUser("ben", "mypass", RoleUser, false)) + require.Nil(t, manager.AddUser("emma", "mypass", RoleUser, false)) + require.Nil(t, manager.ChangeTier("phil", "pro")) + phil, err := manager.User("phil") + require.Nil(t, err) + ben, err := manager.User("ben") + require.Nil(t, err) + emma, err := manager.User("emma") + require.Nil(t, err) + require.Nil(t, manager.MarkUserRemoved(emma)) + users, err := manager.UsersByIDs([]string{phil.ID, ben.ID, emma.ID, "u_doesnotexist"}) + require.Nil(t, err) + require.Len(t, users, 3) + byID := make(map[string]*User) + for _, u := range users { + byID[u.ID] = u + } + require.Equal(t, "phil", byID[phil.ID].Name) + require.Equal(t, "pro", byID[phil.ID].Tier.Code) + require.Equal(t, int64(123), byID[phil.ID].Tier.MessageLimit) + require.Equal(t, "ben", byID[ben.ID].Name) + require.Nil(t, byID[ben.ID].Tier) + require.True(t, byID[emma.ID].Deleted) + users, err = manager.UsersByIDs(nil) + require.Nil(t, err) + require.Empty(t, users) + + // More IDs than one statement takes + ids := make([]string, 0) + for i := 0; i < usersByIDsChunk+5; i++ { + name := fmt.Sprintf("user%d", i) + require.Nil(t, manager.AddUser(name, "mypass", RoleUser, false)) + u, err := manager.User(name) + require.Nil(t, err) + ids = append(ids, u.ID) + } + users, err = manager.UsersByIDs(append(ids, "u_doesnotexist")) + require.Nil(t, err) + require.Len(t, users, usersByIDsChunk+5) + }) +} + func TestStoreUserByToken(t *testing.T) { forEachStoreBackend(t, func(t *testing.T, manager *Manager) { require.Nil(t, manager.AddUser("phil", "mypass", RoleUser, false)) @@ -3723,11 +3775,11 @@ func checkMigrated10To11(t *testing.T, a *Manager, selectVersionQuery string) { var version int require.Nil(t, a.db.QueryRow(selectVersionQuery).Scan(&version)) require.Equal(t, latestSchemaVersion, version) - _, err := a.db.Exec(`INSERT INTO topic (topic, rate_visitor_key, rate_visitor_user_id, seen_at, updated_at) VALUES ('upAAAAAAAAAAAA', 'ip:1.2.3.4', '', 1793380800, 1793380805)`) + _, err := a.db.Exec(`INSERT INTO topic (topic, first_seen_at, subscribed_at, rate_visitor_user_id, rate_visitor_ip, updated_at) VALUES ('upAAAAAAAAAAAA', 1793380700, 1793380800, 'u_1', '1.2.3.4', 1793380805)`) require.Nil(t, err) - var rateVisitorKey string - require.Nil(t, a.db.QueryRow(`SELECT rate_visitor_key FROM topic WHERE topic = 'upAAAAAAAAAAAA' AND updated_at >= 1793380805`).Scan(&rateVisitorKey)) - require.Equal(t, "ip:1.2.3.4", rateVisitorKey) + var userID, ip string + require.Nil(t, a.db.QueryRow(`SELECT rate_visitor_user_id, rate_visitor_ip FROM topic WHERE topic = 'upAAAAAAAAAAAA' AND updated_at >= 1793380805`).Scan(&userID, &ip)) + require.Equal(t, []string{"u_1", "1.2.3.4"}, []string{userID, ip}) } // newTestSQLiteManager opens the SQLite user database and creates a manager on it