From 1b0c0282975a2bf890359b336e09c6b8d5dff88e Mon Sep 17 00:00:00 2001
From: David Florness <david@florness.com>
Date: Tue, 19 Jan 2021 17:59:49 -0500
Subject: [PATCH] Cascade ElectionsMap's Save to Elections

---
 cmd/tallyard/main.go |  4 ++--
 election/election.go |  3 +--
 election/map.go      | 15 +++++++++++++++
 election/msg.go      |  2 +-
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/cmd/tallyard/main.go b/cmd/tallyard/main.go
index 098dcee..e2bb8ea 100644
--- a/cmd/tallyard/main.go
+++ b/cmd/tallyard/main.go
@@ -55,9 +55,9 @@ func main() {
 	if data.Elections == nil {
 		data.Elections = election.NewElectionsMap()
 	}
-	data.Elections.Save = func() {
+	data.Elections.SetSave(func() {
 		data.Save()
-	}
+	})
 
 	syncer := client.Syncer.(*mautrix.DefaultSyncer)
 	syncer.OnEvent(client.Store.(*mautrix.InMemoryStore).UpdateState)
diff --git a/election/election.go b/election/election.go
index 28191ff..d82c549 100644
--- a/election/election.go
+++ b/election/election.go
@@ -26,7 +26,7 @@ type Election struct {
 
 func NewElection(candidates []Candidate, createEventId id.EventID,
 	creationTimestamp int64, creator id.UserID, roomID id.RoomID,
-	title string, save func()) *Election {
+	title string) *Election {
 	return &Election{
 		Candidates:        candidates,
 		CreateEventId:     createEventId,
@@ -34,7 +34,6 @@ func NewElection(candidates []Candidate, createEventId id.EventID,
 		Creator:           creator,
 		Joins:             make(map[id.EventID]*Voter),
 		RoomID:            roomID,
-		Save:              save,
 		Title:             title,
 	}
 }
diff --git a/election/map.go b/election/map.go
index 33ebf32..219bc99 100644
--- a/election/map.go
+++ b/election/map.go
@@ -45,6 +45,18 @@ func (em *ElectionsMap) UnmarshalJSON(b []byte) error {
 	return nil
 }
 
+func (em *ElectionsMap) SetSave(save func ()) {
+	em.Lock()
+	defer em.Unlock()
+	em.Save = save
+	// ensure old elections have Save set
+	for _, el := range em.M {
+		el.Lock()
+		el.Save = save
+		el.Unlock()
+	}
+}
+
 func (em *ElectionsMap) Get(createEventID id.EventID) *Election {
 	em.RLock()
 	defer em.RUnlock()
@@ -76,6 +88,9 @@ func (em *ElectionsMap) SetIfNotExists(createEventID id.EventID, el *Election) {
 }
 
 func (em *ElectionsMap) set(createEventID id.EventID, el *Election) {
+	el.Lock()
+	defer el.Unlock()
+	el.Save = em.Save
 	em.M[createEventID] = el
 	em.insort(createEventID, el)
 }
diff --git a/election/msg.go b/election/msg.go
index 863f2fc..9d87d0b 100644
--- a/election/msg.go
+++ b/election/msg.go
@@ -95,7 +95,7 @@ func OnCreateElectionMessage(evt *event.Event, elections *ElectionsMap) {
 		return
 	}
 	elections.SetIfNotExists(evt.ID, NewElection(content.Candidates, evt.ID,
-		evt.Timestamp, evt.Sender, evt.RoomID, content.Title, elections.Save))
+		evt.Timestamp, evt.Sender, evt.RoomID, content.Title))
 }
 
 func getElection(client *mautrix.Client, roomID id.RoomID, createEventId id.EventID, elections *ElectionsMap) *Election {
-- 
2.38.4