MM-67686: Elasticsearch indexing job progress estimation (#35433)

* Improve job progress estimation with no data

For the Elasticsearch indexing job, we compute the job progress
executing analytics queries to get the total number of posts, channels,
users and files in the database.

Until now, if that call failed, we instead used an estimate. That
estimate is a hardcoded number that is in no way related to the server
data. If that estimate was smaller than the number of already processed
entities, the job progress would show up as larger than 100%.

This commit changes that behaviour by caching the result of the
analytics query:

1. If the analytics query succeeds, we store that value in the job data.
2. If the analytics query fails, we pick a value for the total as
   follows:
   - Use the value previously stored in the job data if available.
   - If not, use the hardcoded estimate.
   - If the hardcoded estimate is smaller than the current count of
     processed entities, use that count instead.

* Add defensive code against division by zero

---------

Co-authored-by: Mattermost Build <build@mattermost.com>

server/enterprise/elasticsearch/common/indexing_job.go

@@ -113,7 +113,14 @@ type IndexingProgress struct {
 func (ip *IndexingProgress) CurrentProgress() int64 {
 	current := ip.DonePostsCount + ip.DoneChannelsCount + ip.DoneUsersCount + ip.DoneFilesCount
 	total := ip.TotalPostsCount + ip.TotalChannelsCount + ip.TotalFilesCount + ip.TotalUsersCount
-	return current * 100 / total
+	if total == 0 {
+		return 100
+	}
+	progress := current * 100 / total
+	if progress > 100 {
+		return 100
+	}
+	return progress
 }
 
 func (ip *IndexingProgress) IsDone(job *model.Job) bool {
@@ -292,6 +299,11 @@ func (worker *IndexerWorker) DoJob(job *model.Job) {
 			job.Data["done_users_count"] = strconv.FormatInt(progress.DoneUsersCount, 10)
 			job.Data["done_files_count"] = strconv.FormatInt(progress.DoneFilesCount, 10)
 
+			job.Data["total_posts_count"] = strconv.FormatInt(progress.TotalPostsCount, 10)
+			job.Data["total_channels_count"] = strconv.FormatInt(progress.TotalChannelsCount, 10)
+			job.Data["total_users_count"] = strconv.FormatInt(progress.TotalUsersCount, 10)
+			job.Data["total_files_count"] = strconv.FormatInt(progress.TotalFilesCount, 10)
+
 			job.Data["start_time"] = strconv.FormatInt(progress.LastEntityTime, 10)
 			job.Data["start_post_id"] = progress.LastPostID
 			job.Data["start_channel_id"] = progress.LastChannelID
@@ -793,25 +805,45 @@ func setStartEntityIDs(progress IndexingProgress, job *model.Job) IndexingProgre
 	return progress
 }
 
+// entityCountFallback returns the best available fallback for a total entity count when the
+// analytics query fails. It prefers a previously stored value from job.Data, then the hardcoded
+// estimate, and finally ensures the result is never less than what has already been processed.
+func entityCountFallback(job *model.Job, dataKey string, estimate, doneCount int64) int64 {
+	fallback := estimate
+	if stored, ok := job.Data[dataKey]; ok {
+		if v, err := strconv.ParseInt(stored, 10, 64); err == nil && v > 0 {
+			fallback = v
+		}
+	}
+	if doneCount > fallback {
+		fallback = doneCount
+	}
+	return fallback
+}
+
 func setEntityCount(logger mlog.LoggerIFace, jobServer *jobs.JobServer, progress IndexingProgress, job *model.Job) IndexingProgress {
 	if job.Data["index_posts"] == "true" {
 		// Counting all posts may fail or timeout when the posts table is large. If this happens, log a warning, but carry
 		// on with the indexing job anyway. The only issue is that the progress % reporting will be inaccurate.
 		if count, err := jobServer.Store.Post().AnalyticsPostCount(&model.PostCountOptions{}); err != nil {
-			logger.Warn("Worker: Failed to fetch total post count for job. An estimated value will be used for progress reporting.", mlog.Int("estimatedPostCount", estimatedPostCount), mlog.Err(err))
-			progress.TotalPostsCount = estimatedPostCount
+			fallback := entityCountFallback(job, "total_posts_count", estimatedPostCount, progress.DonePostsCount)
+			logger.Warn("Worker: Failed to fetch total post count for job. A fallback value will be used for progress reporting.", mlog.Int("fallbackPostCount", fallback), mlog.Err(err))
+			progress.TotalPostsCount = fallback
 		} else {
 			progress.TotalPostsCount = count
+			job.Data["total_posts_count"] = strconv.FormatInt(count, 10)
 		}
 	}
 
 	if job.Data["index_channels"] == "true" {
 		// Same possible fail as above can happen when counting channels
 		if count, err := jobServer.Store.Channel().AnalyticsTypeCount("", ""); err != nil {
-			logger.Warn("Worker: Failed to fetch total channel count for job. An estimated value will be used for progress reporting.", mlog.Int("estimatedChannelCount", estimatedChannelCount), mlog.Err(err))
-			progress.TotalChannelsCount = estimatedChannelCount
+			fallback := entityCountFallback(job, "total_channels_count", estimatedChannelCount, progress.DoneChannelsCount)
+			logger.Warn("Worker: Failed to fetch total channel count for job. A fallback value will be used for progress reporting.", mlog.Int("fallbackChannelCount", fallback), mlog.Err(err))
+			progress.TotalChannelsCount = fallback
 		} else {
 			progress.TotalChannelsCount = count
+			job.Data["total_channels_count"] = strconv.FormatInt(count, 10)
 		}
 	}
 
@@ -821,20 +853,24 @@ func setEntityCount(logger mlog.LoggerIFace, jobServer *jobs.JobServer, progress
 			IncludeBotAccounts: true, // This actually doesn't join with the bots table
 			// since ExcludeRegularUsers is set to false
 		}); err != nil {
-			logger.Warn("Worker: Failed to fetch total user count for job. An estimated value will be used for progress reporting.", mlog.Int("estimatedUserCount", estimatedUserCount), mlog.Err(err))
-			progress.TotalUsersCount = estimatedUserCount
+			fallback := entityCountFallback(job, "total_users_count", estimatedUserCount, progress.DoneUsersCount)
+			logger.Warn("Worker: Failed to fetch total user count for job. A fallback value will be used for progress reporting.", mlog.Int("fallbackUserCount", fallback), mlog.Err(err))
+			progress.TotalUsersCount = fallback
 		} else {
 			progress.TotalUsersCount = count
+			job.Data["total_users_count"] = strconv.FormatInt(count, 10)
 		}
 	}
 
 	if job.Data["index_files"] == "true" {
 		// Same possible fail as above can happen when counting files
 		if count, err := jobServer.Store.FileInfo().CountAll(); err != nil {
-			logger.Warn("Worker: Failed to fetch total files count for job. An estimated value will be used for progress reporting.", mlog.Int("estimatedFilesCount", estimatedFilesCount), mlog.Err(err))
-			progress.TotalFilesCount = estimatedFilesCount
+			fallback := entityCountFallback(job, "total_files_count", estimatedFilesCount, progress.DoneFilesCount)
+			logger.Warn("Worker: Failed to fetch total files count for job. A fallback value will be used for progress reporting.", mlog.Int("fallbackFilesCount", fallback), mlog.Err(err))
+			progress.TotalFilesCount = fallback
 		} else {
 			progress.TotalFilesCount = count
+			job.Data["total_files_count"] = strconv.FormatInt(count, 10)
 		}
 	}
 

server/enterprise/elasticsearch/common/indexing_job_test.go

@@ -4,7 +4,10 @@
 package common
 
 import (
+	"fmt"
 	"io"
+	"maps"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -13,6 +16,7 @@ import (
 
 	"github.com/mattermost/mattermost/server/public/model"
 	"github.com/mattermost/mattermost/server/public/shared/mlog"
+	"github.com/mattermost/mattermost/server/v8/channels/jobs"
 	"github.com/mattermost/mattermost/server/v8/channels/store/storetest/mocks"
 )
 
@@ -63,3 +67,230 @@ func TestBulkIndexChannelsWithDeletedChannels(t *testing.T) {
 	assert.True(t, indexedChannels["ch1"], "Active channel should be indexed")
 	assert.True(t, indexedChannels["ch2"], "Deleted channel should also be indexed")
 }
+
+func TestCurrentProgressCapsAt100(t *testing.T) {
+	tests := []struct {
+		name     string
+		progress IndexingProgress
+		expected int64
+	}{
+		{
+			name: "normal progress",
+			progress: IndexingProgress{
+				DonePostsCount:     50,
+				TotalPostsCount:    100,
+				DoneChannelsCount:  0,
+				TotalChannelsCount: 100,
+				DoneUsersCount:     0,
+				TotalUsersCount:    100,
+				DoneFilesCount:     0,
+				TotalFilesCount:    100,
+			},
+			expected: 12,
+		},
+		{
+			name: "exactly 100%",
+			progress: IndexingProgress{
+				DonePostsCount:     100,
+				TotalPostsCount:    100,
+				DoneChannelsCount:  50,
+				TotalChannelsCount: 50,
+				DoneUsersCount:     10,
+				TotalUsersCount:    10,
+				DoneFilesCount:     5,
+				TotalFilesCount:    5,
+			},
+			expected: 100,
+		},
+		{
+			name:     "all totals zero returns 100",
+			progress: IndexingProgress{},
+			expected: 100,
+		},
+		{
+			name: "done exceeds total, caps at 100",
+			progress: IndexingProgress{
+				DonePostsCount:     20000000,
+				TotalPostsCount:    10000000,
+				DoneChannelsCount:  0,
+				TotalChannelsCount: 100000,
+				DoneUsersCount:     0,
+				TotalUsersCount:    10000,
+				DoneFilesCount:     0,
+				TotalFilesCount:    100000,
+			},
+			expected: 100,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.expected, tt.progress.CurrentProgress())
+		})
+	}
+}
+
+func TestEntityCountFallback(t *testing.T) {
+	tests := []struct {
+		name      string
+		jobData   model.StringMap
+		dataKey   string
+		estimate  int64
+		doneCount int64
+		expected  int64
+	}{
+		{
+			name:      "no stored value, uses estimate",
+			jobData:   model.StringMap{},
+			dataKey:   "total_posts_count",
+			estimate:  10000000,
+			doneCount: 0,
+			expected:  10000000,
+		},
+		{
+			name:      "stored value preferred over estimate",
+			jobData:   model.StringMap{"total_posts_count": "5000000"},
+			dataKey:   "total_posts_count",
+			estimate:  10000000,
+			doneCount: 0,
+			expected:  5000000,
+		},
+		{
+			name:      "doneCount exceeds stored value",
+			jobData:   model.StringMap{"total_posts_count": "5000000"},
+			dataKey:   "total_posts_count",
+			estimate:  10000000,
+			doneCount: 7000000,
+			expected:  7000000,
+		},
+		{
+			name:      "doneCount exceeds estimate when no stored value",
+			jobData:   model.StringMap{},
+			dataKey:   "total_posts_count",
+			estimate:  10000000,
+			doneCount: 15000000,
+			expected:  15000000,
+		},
+		{
+			name:      "invalid stored value falls back to estimate",
+			jobData:   model.StringMap{"total_posts_count": "not-a-number"},
+			dataKey:   "total_posts_count",
+			estimate:  10000000,
+			doneCount: 0,
+			expected:  10000000,
+		},
+		{
+			name:      "zero stored value falls back to estimate",
+			jobData:   model.StringMap{"total_posts_count": "0"},
+			dataKey:   "total_posts_count",
+			estimate:  10000000,
+			doneCount: 0,
+			expected:  10000000,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			job := &model.Job{Data: tt.jobData}
+			result := entityCountFallback(job, tt.dataKey, tt.estimate, tt.doneCount)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestSetEntityCount(t *testing.T) {
+	type entityCountMocks struct {
+		store    *mocks.Store
+		post     *mocks.PostStore
+		channel  *mocks.ChannelStore
+		user     *mocks.UserStore
+		fileInfo *mocks.FileInfoStore
+	}
+
+	setupEntityCountMocks := func() entityCountMocks {
+		m := entityCountMocks{
+			store:    &mocks.Store{},
+			post:     &mocks.PostStore{},
+			channel:  &mocks.ChannelStore{},
+			user:     &mocks.UserStore{},
+			fileInfo: &mocks.FileInfoStore{},
+		}
+		m.store.On("Post").Return(m.post)
+		m.store.On("Channel").Return(m.channel)
+		m.store.On("User").Return(m.user)
+		m.store.On("FileInfo").Return(m.fileInfo)
+		return m
+	}
+
+	allEntitiesEnabled := model.StringMap{
+		"index_posts":    "true",
+		"index_channels": "true",
+		"index_users":    "true",
+		"index_files":    "true",
+	}
+
+	t.Run("stores counts on success", func(t *testing.T) {
+		m := setupEntityCountMocks()
+		m.post.On("AnalyticsPostCount", mock.Anything).Return(int64(500), nil)
+		m.channel.On("AnalyticsTypeCount", "", model.ChannelType("")).Return(int64(200), nil)
+		m.user.On("Count", mock.Anything).Return(int64(50), nil)
+		m.fileInfo.On("CountAll").Return(int64(300), nil)
+
+		job := &model.Job{Data: maps.Clone(allEntitiesEnabled)}
+		progress := setEntityCount(mlog.CreateConsoleTestLogger(t), &jobs.JobServer{Store: m.store}, IndexingProgress{}, job)
+
+		assert.Equal(t, int64(500), progress.TotalPostsCount)
+		assert.Equal(t, int64(200), progress.TotalChannelsCount)
+		assert.Equal(t, int64(50), progress.TotalUsersCount)
+		assert.Equal(t, int64(300), progress.TotalFilesCount)
+
+		assert.Equal(t, "500", job.Data["total_posts_count"])
+		assert.Equal(t, "200", job.Data["total_channels_count"])
+		assert.Equal(t, "50", job.Data["total_users_count"])
+		assert.Equal(t, "300", job.Data["total_files_count"])
+	})
+
+	t.Run("falls back to job data on query failure", func(t *testing.T) {
+		m := setupEntityCountMocks()
+		m.post.On("AnalyticsPostCount", mock.Anything).Return(int64(0), fmt.Errorf("timeout"))
+		m.channel.On("AnalyticsTypeCount", "", model.ChannelType("")).Return(int64(0), fmt.Errorf("timeout"))
+		m.user.On("Count", mock.Anything).Return(int64(0), fmt.Errorf("timeout"))
+		m.fileInfo.On("CountAll").Return(int64(0), fmt.Errorf("timeout"))
+
+		jobData := maps.Clone(allEntitiesEnabled)
+		jobData["total_posts_count"] = "8000000"
+		jobData["total_channels_count"] = "50000"
+		jobData["total_users_count"] = "5000"
+		jobData["total_files_count"] = "75000"
+		job := &model.Job{Data: jobData}
+
+		progress := setEntityCount(mlog.CreateConsoleTestLogger(t), &jobs.JobServer{Store: m.store}, IndexingProgress{}, job)
+
+		assert.Equal(t, int64(8000000), progress.TotalPostsCount)
+		assert.Equal(t, int64(50000), progress.TotalChannelsCount)
+		assert.Equal(t, int64(5000), progress.TotalUsersCount)
+		assert.Equal(t, int64(75000), progress.TotalFilesCount)
+	})
+
+	t.Run("uses max of fallback and done count", func(t *testing.T) {
+		m := setupEntityCountMocks()
+		m.post.On("AnalyticsPostCount", mock.Anything).Return(int64(0), fmt.Errorf("timeout"))
+
+		job := &model.Job{
+			Data: model.StringMap{
+				"index_posts":       "true",
+				"index_channels":    "false",
+				"index_users":       "false",
+				"index_files":       "false",
+				"total_posts_count": strconv.FormatInt(estimatedPostCount, 10),
+			},
+		}
+		inputProgress := IndexingProgress{
+			DonePostsCount: estimatedPostCount + 5000000,
+		}
+
+		progress := setEntityCount(mlog.CreateConsoleTestLogger(t), &jobs.JobServer{Store: m.store}, inputProgress, job)
+
+		assert.Equal(t, int64(estimatedPostCount+5000000), progress.TotalPostsCount)
+	})
+}