Merge pull request #1057 from snimshchikov/main

fix(updates): fix common gaps issues

internal/gen/_template/updates_classifier.tmpl

@@ -0,0 +1,72 @@
+{{ define "updates_classifier" }}
+{{ $pkg := $.Package }}
+{{ template "header" $ }}
+
+func IsPtsUpdate(u UpdateClass) (pts, ptsCount int, ok bool) {
+	switch u := u.(type) {
+{{- range $s := $.Structs }}{{ if and (eq $s.Interface "UpdateClass") 
+		  							  (hasField $s.Fields "Pts" "int") 
+									  (not (contains $s.Name "Channel")) }}
+{{- $ptsCount := or (and (hasField $s.Fields "PtsCount" "int") "u.PtsCount") "0" }}
+	case *{{ $s.Name }}:
+		return u.Pts, {{ $ptsCount }}, true
+{{- end }}{{ end }}
+	}
+
+	return
+}
+
+func IsQtsUpdate(u UpdateClass) (qts int, ok bool) {
+	switch u := u.(type) {
+{{- range $s := $.Structs }}{{ if and (eq $s.Interface "UpdateClass") (hasField $s.Fields "Qts" "int") }}
+	case *{{ $s.Name }}:
+		return u.Qts, true
+{{- end }}{{ end }}
+	}
+
+	return
+}
+
+func IsChannelPtsUpdate(u UpdateClass) (channelID int64, pts, ptsCount int, ok bool, err error) {
+	switch u := u.(type) {
+{{- range $s := $.Structs }}{{ if and (eq $s.Interface "UpdateClass") 
+		  							  (hasField $s.Fields "Pts" "int") 
+									  (contains $s.Name "Channel") }}
+{{- $ptsCount := or (and (hasField $s.Fields "PtsCount" "int") "u.PtsCount") "0" }}
+	case *{{ $s.Name }}:
+{{- if (hasField $s.Fields "ChannelID" "int64") }}
+		return u.ChannelID, u.Pts, {{ $ptsCount }}, true, nil
+{{- else }}
+		channelID, err = extractChannelID(u.Message)
+		return channelID, u.Pts, {{ $ptsCount }}, true, err
+{{- end }}{{ end }}{{ end }}
+	}
+
+	return
+}
+
+func extractChannelID(msg MessageClass) (int64, error) {
+	switch msg := msg.(type) {
+{{- range $s := $.Structs }}{{ if eq $s.Interface "MessageClass" }}
+	case *{{ $s.Name }}:
+{{- range $field := $s.Fields }}{{ if eq $field.Name "PeerID" }}
+{{- if (optionalField $s $field) }}
+		peer, ok := msg.GetPeerID()
+		if !ok {
+			return 0, errors.New("{{ $s.Name }} have no peerID field")
+		}
+{{- else }}
+		peer := msg.PeerID
+{{ end }}{{ end }}{{ end }}
+		if c, ok := peer.(*PeerChannel); ok {
+			return c.ChannelID, nil
+		}
+
+		return 0, errors.New("unexpected peer type")
+{{- end }}{{ end }}
+	default:
+		return 0, errors.New("unexpected MessageClass type")
+	}
+}
+
+{{ end }}

internal/gen/config.go

@@ -12,6 +12,8 @@ type GenerateFlags struct {
 	Server bool
 	// Handlers enables update handler generation.
 	Handlers bool
+	// UpdatesClassifier enables updates classifier generation.
+	UpdatesClassifier bool
 	// GetSet enables getters and setters generation.
 	GetSet bool
 	// Mapping enables mapping helpers generation.
@@ -28,6 +30,7 @@ func (s *GenerateFlags) RegisterFlags(set *flag.FlagSet) {
 	set.BoolVar(&s.Registry, "registry", true, "Enables type ID registry generation")
 	set.BoolVar(&s.Server, "server", false, "Enables experimental server generation")
 	set.BoolVar(&s.Handlers, "handlers", false, "Enables update handler generation")
+	set.BoolVar(&s.UpdatesClassifier, "updates-classifier", true, "Enables updates classifier generation")
 	set.BoolVar(&s.GetSet, "getset", true, "Enables getters and setters generation")
 	set.BoolVar(&s.Mapping, "mapping", false, "Enables mapping helpers generation")
 	set.BoolVar(&s.Slices, "slices", false, "Enables slice helpers generation")

internal/gen/templates.go

@@ -62,7 +62,9 @@ func Funcs() template.FuncMap {
 		"trimSuffix":           strings.TrimSuffix,
 		"hasPrefix":            strings.HasPrefix,
 		"hasSuffix":            strings.HasSuffix,
+		"contains":             strings.Contains,
 		"hasField":             hasField,
+		"optionalField":        optionalField,
 		"mapCollectableFields": mapCollectableFields,
 		"sortableFields":       sortableFields,
 		"generateSliceHelper":  generateSliceHelper,

internal/gen/write_source.go

@@ -178,6 +178,13 @@ func (g *Generator) WriteSource(fs FileSystem, pkgName string, t *template.Templ
 			return err
 		}
 	}
+	if g.generateFlags.UpdatesClassifier && g.hasUpdateClass() {
+		if err := w.Generate("updates_classifier", "tl_updates_classifier_gen.go", config{
+			Structs: g.structs,
+		}); err != nil {
+			return err
+		}
+	}
 
 	sort.SliceStable(g.interfaces, func(i, j int) bool {
 		return g.interfaces[i].Name < g.interfaces[j].Name

telegram/updates/gap_check.go

@@ -10,6 +10,11 @@ const (
 )
 
 func checkGap(localState, remoteState, count int) gapCheckResult {
+	// Temporary fix for handling qts updates gaps.
+	if remoteState == 0 {
+		return gapApply
+	}
+
 	if localState+count == remoteState {
 		return gapApply
 	}

telegram/updates/sort_pts.go

@@ -29,22 +29,22 @@ func (p ptsSorter) Less(i, j int) bool {
 	//
 	// 0) Common updates without PTS.
 	// 1) Common PTS updates
-	// 2) Channel PTS updates (by channelID and by pts)
-	// 3) QTS updates
+	// 2) Common QTS updates
+	// 3) Channel PTS updates (by channelID and by pts)
 	const (
 		other ptsType = iota
 		commonPts
-		channelPts
 		commonQts
+		channelPts
 	)
 	getType := func(u tg.UpdateClass) compare {
-		if pts, ptsCount, ok := isCommonPtsUpdate(u); ok {
+		if pts, ptsCount, ok := tg.IsPtsUpdate(u); ok {
 			return compare{typ: commonPts, ptsDiff: pts - ptsCount}
 		}
-		if channelID, pts, ptsCount, ok, _ := isChannelPtsUpdate(u); ok {
+		if channelID, pts, ptsCount, ok, _ := tg.IsChannelPtsUpdate(u); ok {
 			return compare{typ: channelPts, channelID: channelID, ptsDiff: pts - ptsCount}
 		}
-		if qts, ok := isCommonQtsUpdate(u); ok {
+		if qts, ok := tg.IsQtsUpdate(u); ok {
 			return compare{typ: commonQts, ptsDiff: qts}
 		}
 		return compare{typ: other}

telegram/updates/sort_pts_test.go

@@ -65,12 +65,12 @@ func Test_sortUpdatesByPts(t *testing.T) {
 			[]tg.UpdateClass{
 				newMessage(25),
 				newMessage(26),
-				channelNewMessage(25, 1),
-				channelReadInbox(26, 1),
-				channelNewMessage(25, 2),
 				encryptedNewMessage(25),
 				encryptedNewMessage(26),
 				encryptedNewMessage(27),
+				channelNewMessage(25, 1),
+				channelReadInbox(26, 1),
+				channelNewMessage(25, 2),
 			},
 		},
 		{

telegram/updates/state_apply.go

@@ -67,15 +67,15 @@ func (s *internalState) applyCombined(ctx context.Context, comb *tg.UpdatesCombi
 			continue
 		}
 
-		if pts, ptsCount, ok := isCommonPtsUpdate(u); ok {
+		if pts, ptsCount, ok := tg.IsPtsUpdate(u); ok {
 			if err := s.handlePts(ctx, pts, ptsCount, u, ents); err != nil {
 				return false, err
 			}
 
 			continue
 		}
 
-		if channelID, pts, ptsCount, ok, err := isChannelPtsUpdate(u); ok {
+		if channelID, pts, ptsCount, ok, err := tg.IsChannelPtsUpdate(u); ok {
 			if err != nil {
 				s.log.Debug("Invalid channel update", zap.Error(err), zap.Any("update", u))
 				continue
@@ -91,7 +91,7 @@ func (s *internalState) applyCombined(ctx context.Context, comb *tg.UpdatesCombi
 			continue
 		}
 
-		if qts, ok := isCommonQtsUpdate(u); ok {
+		if qts, ok := tg.IsQtsUpdate(u); ok {
 			if err := s.handleQts(ctx, qts, u, ents); err != nil {
 				return false, err
 			}
@@ -187,6 +187,11 @@ func (s *internalState) applyQts(ctx context.Context, state int, updates []updat
 		s.log.Error("Handle updates error", zap.Error(err))
 	}
 
+	// Don't set qts if it's 0, because it means that we are apllying gaps updates
+	if state == 0 {
+		return nil
+	}
+
 	if err := s.storage.SetQts(ctx, s.selfID, state); err != nil {
 		s.log.Error("SetQts error", zap.Error(err))
 	}

telegram/updates/state_channel.go

@@ -135,7 +135,7 @@ func (s *channelState) handleUpdate(ctx context.Context, u tg.UpdateClass, ents
 		return s.handleTooLong(ctx, long)
 	}
 
-	channelID, pts, ptsCount, ok, err := isChannelPtsUpdate(u)
+	channelID, pts, ptsCount, ok, err := tg.IsChannelPtsUpdate(u)
 	if err != nil {
 		return errors.Wrap(err, "invalid update")
 	}