Use timers to call removeUser after timeout

src/github.com/matrix-org/dendrite/typingserver/cache/cache.go

@@ -19,8 +19,8 @@ import (
 
 var defaultTypingTimeout = 10 * time.Second
 
-// userSet is a map of user IDs to their time of expiry.
-type userSet map[string]time.Time
+// userSet is a map of user IDs to a timer, timer fires at expiry.
+type userSet map[string]*time.Timer
 
 // TypingCache maintains a list of users typing in each room.
 type TypingCache struct {
@@ -53,14 +53,14 @@ func (t *TypingCache) GetTypingUsers(roomID string) (users []string) {
 // if expire is nil, defaultTypingTimeout is assumed.
 func (t *TypingCache) AddTypingUser(userID, roomID string, expire *time.Time) {
 	expireTime := getExpireTime(expire)
-	if time.Until(expireTime) > 0 {
-		t.addUser(userID, roomID, expireTime)
-		t.removeUserAfterTime(userID, roomID, expireTime)
+	if until := time.Until(expireTime); until > 0 {
+		timer := time.AfterFunc(until, t.timeoutCallback(userID, roomID))
+		t.addUser(userID, roomID, timer)
 	}
 }
 
-// addUser with mutex lock.
-func (t *TypingCache) addUser(userID, roomID string, expireTime time.Time) {
+// addUser with mutex lock & replace the previous timer.
+func (t *TypingCache) addUser(userID, roomID string, expiryTimer *time.Timer) {
 	t.Lock()
 	defer t.Unlock()
 
@@ -68,24 +68,34 @@ func (t *TypingCache) addUser(userID, roomID string, expireTime time.Time) {
 		t.data[roomID] = make(userSet)
 	}
 
-	t.data[roomID][userID] = expireTime
+	// Stop the timer to cancel the call to timeoutCallback
+	if timer, ok := t.data[roomID][userID]; ok {
+		// It may happen that at this stage timer fires but now we have a lock on t.
+		// Hence the execution of timeoutCallback will happen after we unlock.
+		// So we may lose a typing state, though this event is highly unlikely.
+		// This can be mitigated by keeping another time.Time in the map and check against it
+		// before removing. This however is not required in most practical scenario.
+		timer.Stop()
 	}
 
-// Creates a go routine which removes the user after expireTime has elapsed,
-// only if the expiration is not updated to a later time in cache.
-func (t *TypingCache) removeUserAfterTime(userID, roomID string, expireTime time.Time) {
-	go func() {
-		time.Sleep(time.Until(expireTime))
-		t.removeUserIfExpired(userID, roomID)
-	}()
+	t.data[roomID][userID] = expiryTimer
 }
 
-// removeUserIfExpired with mutex lock.
-func (t *TypingCache) removeUserIfExpired(userID, roomID string) {
+// Returns a function which is called after timeout happens.
+// This removes the user.
+func (t *TypingCache) timeoutCallback(userID, roomID string) func() {
+	return func() {
+		t.removeUser(userID, roomID)
+	}
+}
+
+// removeUser with mutex lock & stop the timer.
+func (t *TypingCache) removeUser(userID, roomID string) {
 	t.Lock()
 	defer t.Unlock()
 
-	if time.Until(t.data[roomID][userID]) <= 0 {
+	if timer, ok := t.data[roomID][userID]; ok {
+		timer.Stop()
 		delete(t.data[roomID], userID)
 	}
 }

src/github.com/matrix-org/dendrite/typingserver/cache/cache_test.go

@@ -19,8 +19,6 @@ import (
 	"github.com/matrix-org/dendrite/common/test"
 )
 
-const longInterval = time.Hour
-
 func TestTypingCache(t *testing.T) {
 	tCache := NewTypingCache()
 	if tCache == nil {
@@ -34,14 +32,10 @@ func TestTypingCache(t *testing.T) {
 	t.Run("GetTypingUsers", func(t *testing.T) {
 		testGetTypingUsers(t, tCache)
 	})
-
-	t.Run("removeUserIfExpired", func(t *testing.T) {
-		testRemoveUserIfExpired(t, tCache)
-	})
 }
 
 func testAddTypingUser(t *testing.T, tCache *TypingCache) {
-	timeAfterLongInterval := time.Now().Add(longInterval)
+	present := time.Now()
 	tests := []struct {
 		userID string
 		roomID string
@@ -51,8 +45,8 @@ func testAddTypingUser(t *testing.T, tCache *TypingCache) {
 		{"user2", "room1", nil},
 		{"user3", "room1", nil},
 		{"user4", "room1", nil},
-		// removeUserIfExpired should not remove the user before expiration time.
-		{"user1", "room2", &timeAfterLongInterval},
+		//typing state with past expireTime should not take effect or removed.
+		{"user1", "room2", &present},
 	}
 
 	for _, tt := range tests {
@@ -66,7 +60,7 @@ func testGetTypingUsers(t *testing.T, tCache *TypingCache) {
 		wantUsers []string
 	}{
 		{"room1", []string{"user1", "user2", "user3", "user4"}},
-		{"room2", []string{"user1"}},
+		{"room2", []string{}},
 	}
 
 	for _, tt := range tests {
@@ -76,20 +70,3 @@ func testGetTypingUsers(t *testing.T, tCache *TypingCache) {
 		}
 	}
 }
-
-func testRemoveUserIfExpired(t *testing.T, tCache *TypingCache) {
-	tests := []struct {
-		roomID    string
-		userID    string
-		wantUsers []string
-	}{
-		{"room2", "user1", []string{"user1"}},
-	}
-
-	for _, tt := range tests {
-		tCache.removeUserIfExpired(tt.userID, tt.roomID)
-		if gotUsers := tCache.GetTypingUsers(tt.roomID); !test.UnsortedStringSliceEqual(gotUsers, tt.wantUsers) {
-			t.Errorf("TypingCache.GetTypingUsers(%s) = %v, want %v", tt.roomID, gotUsers, tt.wantUsers)
-		}
-	}
-}