Handle unsubscription well

.idea/vcs.xml

@@ -9,6 +9,11 @@
       </inspection_tool>
     </profile>
   </component>
+  <component name="GitSharedSettings">
+    <option name="FORCE_PUSH_PROHIBITED_PATTERNS">
+      <list />
+    </option>
+  </component>
   <component name="VcsDirectoryMappings">
     <mapping directory="$PROJECT_DIR$" vcs="Git" />
   </component>

main.go

@@ -257,7 +257,7 @@ func handler(ctx context.Context, b *tgBot.Bot, update *models.Update, subs *sub
 						},
 					}
 				} else {
-					// TODO: Update message to contain unsubscribe button
+					log.Printf("DEBUG: Subscribed: chatID %d, trainNumber %s, date %s, groupIndex %d", update.CallbackQuery.Message.Chat.ID, trainNumber, date.Format("2006-01-02"), groupIndex)
 					response = &handlers.HandlerResponse{
 						CallbackAnswer: &tgBot.AnswerCallbackQueryParams{
 							Text: fmt.Sprintf("Subscribed successfully!"),
@@ -287,7 +287,7 @@ func handler(ctx context.Context, b *tgBot.Bot, update *models.Update, subs *sub
 						},
 					}
 				} else {
-					// TODO: Update message to contain unsubscribe button
+					log.Printf("DEBUG: Unsubscribed: chatID %d, trainNumber %s, date %s, groupIndex %d", update.CallbackQuery.Message.Chat.ID, trainNumber, date.Format("2006-01-02"), groupIndex)
 					response = &handlers.HandlerResponse{
 						CallbackAnswer: &tgBot.AnswerCallbackQueryParams{
 							Text: fmt.Sprintf("Unsubscribed successfully!"),
@@ -301,6 +301,9 @@ func handler(ctx context.Context, b *tgBot.Bot, update *models.Update, subs *sub
 						},
 					}
 				}
+
+			default:
+				log.Printf("WARN : Unknown callback query method: %s", splitted[0])
 			}
 		}
 	}

pkg/handlers/findTrain.go

@@ -46,6 +46,11 @@ func HandleTrainNumberCommand(ctx context.Context, trainNumber string, date time
 			Message: &bot.SendMessageParams{
 				Text: fmt.Sprintf("The train %s was not found.", trainNumber),
 			},
+			ShouldUnsubscribe: func() bool {
+				now := time.Now().In(utils.Location)
+				midnightYesterday := time.Date(now.Year(), now.Month(), now.Day()-1, 0, 0, 0, 0, utils.Location)
+				return date.Before(midnightYesterday)
+			}(),
 		}, false
 	case errors.Is(err, api.ServerError):
 		log.Printf("ERROR: In handle train number: %s", err.Error())
@@ -53,6 +58,11 @@ func HandleTrainNumberCommand(ctx context.Context, trainNumber string, date time
 			Message: &bot.SendMessageParams{
 				Text: fmt.Sprintf("Unknown server error when searching for train %s.", trainNumber),
 			},
+			ShouldUnsubscribe: func() bool {
+				now := time.Now().In(utils.Location)
+				midnightYesterday := time.Date(now.Year(), now.Month(), now.Day()-1, 0, 0, 0, 0, utils.Location)
+				return date.Before(midnightYesterday)
+			}(),
 		}, false
 	default:
 		log.Printf("ERROR: In handle train number: %s", err.Error())
@@ -63,6 +73,32 @@ func HandleTrainNumberCommand(ctx context.Context, trainNumber string, date time
 		groupIndex = 0
 	}
 
+	shouldUnsubscribe := func() bool {
+		if len(trainData.Groups) <= groupIndex {
+			groupIndex = 0
+		}
+		now := time.Now().In(utils.Location)
+		midnightYesterday := time.Date(now.Year(), now.Month(), now.Day()-1, 0, 0, 0, 0, utils.Location)
+		lastStation := trainData.Groups[groupIndex].
+			Stations[len(trainData.Groups[groupIndex].Stations)-1]
+		if now.After(lastStation.Arrival.
+			ScheduleTime.Add(time.Hour * 6)) {
+			return true
+		}
+		if trainData.Groups[groupIndex].
+			Status != nil && trainData.Groups[groupIndex].
+			Status.Station == lastStation.Name &&
+			trainData.Groups[groupIndex].Status.
+				State == "arrival" {
+			return true
+		}
+		if date.Before(midnightYesterday) {
+			return true
+		}
+
+		return false
+	}()
+
 	message := bot.SendMessageParams{}
 	if groupIndex == -1 {
 		message.Text = fmt.Sprintf("Train %s %s contains multiple groups. Please choose one.", trainData.Rank, trainData.Number)
@@ -136,7 +172,9 @@ func HandleTrainNumberCommand(ctx context.Context, trainNumber string, date time
 			},
 		}
 		buttonKind := TrainInfoResponseButtonIncludeSub
-		if isSubscribed {
+		if shouldUnsubscribe {
+			buttonKind = TrainInfoResponseButtonExcludeSub
+		} else if isSubscribed {
 			buttonKind = TrainInfoResponseButtonIncludeUnsub
 		}
 		message.ReplyMarkup = GetTrainNumberCommandResponseButtons(trainData.Number, group.Stations[0].Departure.ScheduleTime, groupIndex, buttonKind)
@@ -154,6 +192,7 @@ func HandleTrainNumberCommand(ctx context.Context, trainNumber string, date time
 
 	return &HandlerResponse{
 		Message:           &message,
+		ShouldUnsubscribe: shouldUnsubscribe,
 	}, true
 }
 

pkg/handlers/response.go

@@ -8,6 +8,7 @@ type HandlerResponse struct {
 	CallbackAnswer          *bot.AnswerCallbackQueryParams
 	MessageEdits            []*bot.EditMessageTextParams
 	MessageMarkupEdits      []*bot.EditMessageReplyMarkupParams
+	ShouldUnsubscribe       bool
 	Injected                struct {
 		ChatId    int64
 		MessageId int

pkg/subscriptions/subscriptions.go

@@ -98,15 +98,14 @@ func (sub *Subscriptions) DeleteSubscription(chatId int64, messageId int) (*SubD
 			break
 		}
 	}
-	var result *SubData
+	var result SubData
 	if deleteIndex != -1 {
-		result = &SubData{}
+		result = datas[deleteIndex]
-		*result = datas[deleteIndex]
 		datas[deleteIndex] = datas[len(datas)-1]
 		datas = datas[:len(datas)-1]
 
 		_, err := database.WriteDB(func(db *gorm.DB) (*gorm.DB, error) {
-			db.Delete(result)
+			db.Delete(&result)
 			return db, db.Error
 		})
 		if err != nil {
@@ -120,7 +119,7 @@ func (sub *Subscriptions) DeleteSubscription(chatId int64, messageId int) (*SubD
 	} else {
 		sub.data[chatId] = datas
 	}
-	return result, nil
+	return &result, nil
 }
 
 func (sub *Subscriptions) CheckSubscriptions(ctx context.Context) {
@@ -142,20 +141,29 @@ type workerData struct {
 	data  SubData
 }
 
+type unsubscribe struct {
+	chatId    int64
+	messageId int
+}
+
+type workerResponseData struct {
+	unsubscribe *unsubscribe
+}
+
 func (sub *Subscriptions) executeChecks(ctx context.Context) {
 	sub.mutex.RLock()
-	defer sub.mutex.RUnlock()
 
 	// Only allow 8 concurrent requests
 	// TODO: Make configurable instead of hardcoded
 	workerCount := 8
 	workerChan := make(chan workerData, workerCount)
-	wg := &sync.WaitGroup{}
+	responseChan := make(chan *workerResponseData, workerCount)
+	defer close(responseChan)
 	for i := 0; i < workerCount; i++ {
-		wg.Add(1)
+		go checkWorker(ctx, workerChan, responseChan)
-		go checkWorker(ctx, workerChan, wg)
 	}
 
+	go func() {
 		for _, datas := range sub.data {
 			for i := range datas {
 				workerChan <- workerData{
@@ -165,12 +173,42 @@ func (sub *Subscriptions) executeChecks(ctx context.Context) {
 			}
 		}
 		close(workerChan)
-	wg.Wait()
+	}()
+
+	responses := make([]*workerResponseData, 0, len(sub.data))
+
+	for _, datas := range sub.data {
+		for range datas {
+			if resp := <-responseChan; resp != nil && resp.unsubscribe != nil {
+				responses = append(responses, resp)
+			}
+		}
 	}
 
-func checkWorker(ctx context.Context, workerChan <-chan workerData, wg *sync.WaitGroup) {
+	sub.mutex.RUnlock()
-	defer wg.Done()
+
+	for i := range responses {
+		if responses[i].unsubscribe != nil {
+			// Ignore error since this is optional optimisation
+			deletedSub, err := sub.DeleteSubscription(responses[i].unsubscribe.chatId, responses[i].unsubscribe.messageId)
+			if err == nil && deletedSub != nil {
+				_, _ = sub.tgBot.EditMessageReplyMarkup(ctx, &bot.EditMessageReplyMarkupParams{
+					ChatID:      responses[i].unsubscribe.chatId,
+					MessageID:   responses[i].unsubscribe.messageId,
+					ReplyMarkup: handlers.GetTrainNumberCommandResponseButtons(deletedSub.TrainNumber, deletedSub.Date, deletedSub.GroupIndex, handlers.TrainInfoResponseButtonExcludeSub),
+				})
+			}
+		}
+	}
+}
+
+func checkWorker(ctx context.Context, workerChan <-chan workerData, responseChan chan<- *workerResponseData) {
 	for wData := range workerChan {
+		func() {
+			var response *workerResponseData
+			defer func() {
+				responseChan <- response
+			}()
 			data := wData.data
 			log.Printf("DEBUG: Timer tick, update for chat %d, train %s, date %s, group %d", data.ChatId, data.TrainNumber, data.Date.Format("2006-01-02"), data.GroupIndex)
 
@@ -179,6 +217,14 @@ func checkWorker(ctx context.Context, workerChan <-chan workerData, wg *sync.Wai
 			if !ok || resp == nil || resp.Message == nil {
 				// Silently discard update errors
 				log.Printf("DEBUG: Error when updating chat %d, train %s, date %s, group %d", data.ChatId, data.TrainNumber, data.Date.Format("2006-01-02"), data.GroupIndex)
+				if resp != nil && resp.ShouldUnsubscribe {
+					response = &workerResponseData{
+						unsubscribe: &unsubscribe{
+							chatId:    data.ChatId,
+							messageId: data.MessageId,
+						},
+					}
+				}
 				return
 			}
 
@@ -191,5 +237,14 @@ func checkWorker(ctx context.Context, workerChan <-chan workerData, wg *sync.Wai
 				DisableWebPagePreview: resp.Message.DisableWebPagePreview,
 				ReplyMarkup:           resp.Message.ReplyMarkup,
 			})
+
+			response = &workerResponseData{}
+			if resp.ShouldUnsubscribe {
+				response.unsubscribe = &unsubscribe{
+					chatId:    data.ChatId,
+					messageId: data.MessageId,
+				}
+			}
+		}()
 	}
 }