[feature] add support for polls + receiving federated status edits (#2330)

3 files changed, 133 insertions, 10 deletions

diff --git a/internal/cache/cache.go b/internal/cache/cache.go
index 777088f07..dabf151ff 100644
--- a/internal/cache/cache.go
+++ b/internal/cache/cache.go
@@ -183,6 +183,22 @@ func (c *Caches) setuphooks() {
 		}
 	})
 
+	c.GTS.Poll().SetInvalidateCallback(func(poll *gtsmodel.Poll) {
+		// Invalidate all cached votes of this poll.
+		c.GTS.PollVote().Invalidate("PollID", poll.ID)
+
+		// Invalidate cache of poll vote IDs.
+		c.GTS.PollVoteIDs().Invalidate(poll.ID)
+	})
+
+	c.GTS.PollVote().SetInvalidateCallback(func(vote *gtsmodel.PollVote) {
+		// Invalidate cached poll (contains no. votes).
+		c.GTS.Poll().Invalidate("ID", vote.PollID)
+
+		// Invalidate cache of poll vote IDs.
+		c.GTS.PollVoteIDs().Invalidate(vote.PollID)
+	})
+
 	c.GTS.Status().SetInvalidateCallback(func(status *gtsmodel.Status) {
 		// Invalidate status ID cached visibility.
 		c.Visibility.Invalidate("ItemID", status.ID)
@@ -206,6 +222,11 @@ func (c *Caches) setuphooks() {
 			// Invalidate in reply to ID list of original status.
 			c.GTS.InReplyToIDs().Invalidate(status.InReplyToID)
 		}
+
+		if status.PollID != "" {
+			// Invalidate cache of attached poll ID.
+			c.GTS.Poll().Invalidate("ID", status.PollID)
+		}
 	})
 
 	c.GTS.StatusFave().SetInvalidateCallback(func(fave *gtsmodel.StatusFave) {
@@ -244,6 +265,7 @@ func (c *Caches) Sweep(threshold float64) {
 	c.GTS.Media().Trim(threshold)
 	c.GTS.Mention().Trim(threshold)
 	c.GTS.Notification().Trim(threshold)
+	c.GTS.Poll().Trim(threshold)
 	c.GTS.Report().Trim(threshold)
 	c.GTS.Status().Trim(threshold)
 	c.GTS.StatusFave().Trim(threshold)
diff --git a/internal/cache/gts.go b/internal/cache/gts.go
index d96b7267f..de49ad127 100644
--- a/internal/cache/gts.go
+++ b/internal/cache/gts.go
@@ -52,6 +52,9 @@ type GTSCaches struct {
 	media            *result.Cache[*gtsmodel.MediaAttachment]
 	mention          *result.Cache[*gtsmodel.Mention]
 	notification     *result.Cache[*gtsmodel.Notification]
+	poll             *result.Cache[*gtsmodel.Poll]
+	pollVote         *result.Cache[*gtsmodel.PollVote]
+	pollVoteIDs      *SliceCache[string]
 	report           *result.Cache[*gtsmodel.Report]
 	status           *result.Cache[*gtsmodel.Status]
 	statusFave       *result.Cache[*gtsmodel.StatusFave]
@@ -90,6 +93,9 @@ func (c *GTSCaches) Init() {
 	c.initMedia()
 	c.initMention()
 	c.initNotification()
+	c.initPoll()
+	c.initPollVote()
+	c.initPollVoteIDs()
 	c.initReport()
 	c.initStatus()
 	c.initStatusFave()
@@ -231,6 +237,21 @@ func (c *GTSCaches) Notification() *result.Cache[*gtsmodel.Notification] {
 	return c.notification
 }
 
+// Poll provides access to the gtsmodel Poll database cache.
+func (c *GTSCaches) Poll() *result.Cache[*gtsmodel.Poll] {
+	return c.poll
+}
+
+// PollVote provides access to the gtsmodel PollVote database cache.
+func (c *GTSCaches) PollVote() *result.Cache[*gtsmodel.PollVote] {
+	return c.pollVote
+}
+
+// PollVoteIDs provides access to the poll vote IDs list database cache.
+func (c *GTSCaches) PollVoteIDs() *SliceCache[string] {
+	return c.pollVoteIDs
+}
+
 // Report provides access to the gtsmodel Report database cache.
 func (c *GTSCaches) Report() *result.Cache[*gtsmodel.Report] {
 	return c.report
@@ -246,26 +267,26 @@ func (c *GTSCaches) StatusFave() *result.Cache[*gtsmodel.StatusFave] {
 	return c.statusFave
 }
 
-// Tag provides access to the gtsmodel Tag database cache.
-func (c *GTSCaches) Tag() *result.Cache[*gtsmodel.Tag] {
-	return c.tag
-}
-
-// ThreadMute provides access to the gtsmodel ThreadMute database cache.
-func (c *GTSCaches) ThreadMute() *result.Cache[*gtsmodel.ThreadMute] {
-	return c.threadMute
-}
-
 // StatusFaveIDs provides access to the status fave IDs list database cache.
 func (c *GTSCaches) StatusFaveIDs() *SliceCache[string] {
 	return c.statusFaveIDs
 }
 
+// Tag provides access to the gtsmodel Tag database cache.
+func (c *GTSCaches) Tag() *result.Cache[*gtsmodel.Tag] {
+	return c.tag
+}
+
 // Tombstone provides access to the gtsmodel Tombstone database cache.
 func (c *GTSCaches) Tombstone() *result.Cache[*gtsmodel.Tombstone] {
 	return c.tombstone
 }
 
+// ThreadMute provides access to the gtsmodel ThreadMute database cache.
+func (c *GTSCaches) ThreadMute() *result.Cache[*gtsmodel.ThreadMute] {
+	return c.threadMute
+}
+
 // User provides access to the gtsmodel User database cache.
 func (c *GTSCaches) User() *result.Cache[*gtsmodel.User] {
 	return c.user
@@ -685,6 +706,63 @@ func (c *GTSCaches) initNotification() {
 	c.notification.IgnoreErrors(ignoreErrors)
 }
 
+func (c *GTSCaches) initPoll() {
+	// Calculate maximum cache size.
+	cap := calculateResultCacheMax(
+		sizeofPoll(), // model in-mem size.
+		config.GetCachePollMemRatio(),
+	)
+
+	log.Infof(nil, "cache size = %d", cap)
+
+	c.poll = result.New([]result.Lookup{
+		{Name: "ID"},
+		{Name: "StatusID"},
+	}, func(p1 *gtsmodel.Poll) *gtsmodel.Poll {
+		p2 := new(gtsmodel.Poll)
+		*p2 = *p1
+		return p2
+	}, cap)
+
+	c.poll.IgnoreErrors(ignoreErrors)
+}
+
+func (c *GTSCaches) initPollVote() {
+	// Calculate maximum cache size.
+	cap := calculateResultCacheMax(
+		sizeofPollVote(), // model in-mem size.
+		config.GetCachePollVoteMemRatio(),
+	)
+
+	log.Infof(nil, "cache size = %d", cap)
+
+	c.pollVote = result.New([]result.Lookup{
+		{Name: "ID"},
+		{Name: "PollID.AccountID"},
+		{Name: "PollID", Multi: true},
+	}, func(v1 *gtsmodel.PollVote) *gtsmodel.PollVote {
+		v2 := new(gtsmodel.PollVote)
+		*v2 = *v1
+		return v2
+	}, cap)
+
+	c.pollVote.IgnoreErrors(ignoreErrors)
+}
+
+func (c *GTSCaches) initPollVoteIDs() {
+	// Calculate maximum cache size.
+	cap := calculateSliceCacheMax(
+		config.GetCachePollVoteIDsMemRatio(),
+	)
+
+	log.Infof(nil, "cache size = %d", cap)
+
+	c.pollVoteIDs = &SliceCache[string]{Cache: simple.New[string, []string](
+		0,
+		cap,
+	)}
+}
+
 func (c *GTSCaches) initReport() {
 	// Calculate maximum cache size.
 	cap := calculateResultCacheMax(
diff --git a/internal/cache/size.go b/internal/cache/size.go
index 3e77c68ac..62fb31469 100644
--- a/internal/cache/size.go
+++ b/internal/cache/size.go
@@ -189,6 +189,8 @@ func totalOfRatios() float64 {
 		config.GetCacheMediaMemRatio() +
 		config.GetCacheMentionMemRatio() +
 		config.GetCacheNotificationMemRatio() +
+		config.GetCachePollMemRatio() +
+		config.GetCachePollVoteMemRatio() +
 		config.GetCacheReportMemRatio() +
 		config.GetCacheStatusMemRatio() +
 		config.GetCacheStatusFaveMemRatio() +
@@ -438,6 +440,27 @@ func sizeofNotification() uintptr {
 	}))
 }
 
+func sizeofPoll() uintptr {
+	return uintptr(size.Of(&gtsmodel.Poll{
+		ID:         exampleID,
+		Multiple:   func() *bool { ok := false; return &ok }(),
+		HideCounts: func() *bool { ok := false; return &ok }(),
+		Options:    []string{exampleTextSmall, exampleTextSmall, exampleTextSmall, exampleTextSmall},
+		StatusID:   exampleID,
+		ExpiresAt:  exampleTime,
+	}))
+}
+
+func sizeofPollVote() uintptr {
+	return uintptr(size.Of(&gtsmodel.PollVote{
+		ID:        exampleID,
+		Choices:   []int{69, 420, 1337},
+		AccountID: exampleID,
+		PollID:    exampleID,
+		CreatedAt: exampleTime,
+	}))
+}
+
 func sizeofReport() uintptr {
 	return uintptr(size.Of(&gtsmodel.Report{
 		ID:                     exampleID,