upstream-mattermost/server/enterprise/elasticsearch/opensearch/bulk_test.go

// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
// See LICENSE.enterprise for license information.

package opensearch

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/elastic/go-elasticsearch/v8/typedapi/types"
	"github.com/mattermost/mattermost/server/public/model"
	"github.com/mattermost/mattermost/server/public/shared/mlog"
	"github.com/mattermost/mattermost/server/v8/channels/api4"
	"github.com/mattermost/mattermost/server/v8/enterprise/elasticsearch/common"
	"github.com/opensearch-project/opensearch-go/v4"
	"github.com/opensearch-project/opensearch-go/v4/opensearchapi"
	"github.com/stretchr/testify/require"
)

// setupBulkClient creates a test bulk client with common setup
func setupBulkClient(t *testing.T, flushBytes int, flushNumReqs int, flushInterval time.Duration) *Bulk {
	th := api4.SetupEnterprise(t)

	if os.Getenv("IS_CI") == "true" {
		os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://opensearch:9201")
		os.Setenv("MM_ELASTICSEARCHSETTINGS_BACKEND", "opensearch")
	}

	th.App.UpdateConfig(func(cfg *model.Config) {
		if os.Getenv("IS_CI") == "true" {
			*cfg.ElasticsearchSettings.ConnectionURL = "http://opensearch:9201"
		} else {
			*cfg.ElasticsearchSettings.ConnectionURL = "http://localhost:9201"
		}
		*cfg.ElasticsearchSettings.Backend = model.ElasticsearchSettingsOSBackend
		*cfg.ElasticsearchSettings.EnableIndexing = true
		*cfg.ElasticsearchSettings.EnableSearching = true
		*cfg.ElasticsearchSettings.EnableAutocomplete = true
	})

	client := createTestClient(t, th.Context, th.App.Config(), th.App.FileBackend())
	bulk := NewBulk(
		common.BulkSettings{
			FlushBytes:    flushBytes,
			FlushInterval: flushInterval,
			FlushNumReqs:  flushNumReqs,
		},
		client,
		time.Duration(*th.App.Config().ElasticsearchSettings.RequestTimeoutSeconds)*time.Second,
		th.Server.Platform().Log())

	return bulk
}

// createTestPost creates a test post for indexing
func createTestPost(t *testing.T, message string) *common.ESPost {
	post, err := common.ESPostFromPost(&model.Post{
		Id:      model.NewId(),
		Message: message,
	}, "myteam", "O", true)
	require.NoError(t, err)
	return post
}

func TestBulkProcessor(t *testing.T) {
	bulk := setupBulkClient(t, 0, 10, 0)
	defer func() {
		if os.Getenv("IS_CI") == "true" {
			os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
			os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
		}
	}()

	post := createTestPost(t, "hello world")

	err := bulk.IndexOp(&types.IndexOperation{
		Index_: new("myindex"),
		Id_:    new(post.Id),
	}, post)
	require.NoError(t, err)

	require.Equal(t, 1, bulk.pendingRequests)

	err = bulk.Stop()
	require.NoError(t, err)

	require.Equal(t, 0, bulk.pendingRequests)
}

func TestNewBulk(t *testing.T) {
	bulk := setupBulkClient(t, 1024, 10, 0)
	defer func() {
		if os.Getenv("IS_CI") == "true" {
			os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
			os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
		}
	}()

	t.Run("creates bulk client without periodic flusher", func(t *testing.T) {
		require.NotNil(t, bulk)
		require.NotNil(t, bulk.client)
		require.NotNil(t, bulk.logger)
		require.NotNil(t, bulk.buf)
		require.Equal(t, 0, bulk.pendingRequests)
		require.Equal(t, 1024, bulk.bulkSettings.FlushBytes)
		require.Equal(t, 10, bulk.bulkSettings.FlushNumReqs)
	})

	t.Run("creates bulk client with periodic flusher", func(t *testing.T) {
		bulkWithTimer := setupBulkClient(t, 1024, 10, 100*time.Millisecond)
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		require.NotNil(t, bulkWithTimer)
		require.Equal(t, 100*time.Millisecond, bulkWithTimer.bulkSettings.FlushInterval)

		err := bulkWithTimer.Stop()
		require.NoError(t, err)
	})

	err := bulk.Stop()
	require.NoError(t, err)
}

func TestIndexOp(t *testing.T) {
	bulk := setupBulkClient(t, 0, 10, 0)
	defer func() {
		if os.Getenv("IS_CI") == "true" {
			os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
			os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
		}
	}()

	t.Run("single index operation with struct", func(t *testing.T) {
		post := createTestPost(t, "test message")

		err := bulk.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(post.Id),
		}, post)
		require.NoError(t, err)
		require.Equal(t, 1, bulk.pendingRequests)

		// Verify buffer has content
		require.Greater(t, bulk.buf.Len(), 0)
	})

	t.Run("index operation with []byte", func(t *testing.T) {
		initialRequests := bulk.pendingRequests
		docId := model.NewId()
		data := []byte(`{"message": "test byte slice"}`)

		err := bulk.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(docId),
		}, data)
		require.NoError(t, err)
		require.Equal(t, initialRequests+1, bulk.pendingRequests)
	})

	t.Run("index operation with json.RawMessage", func(t *testing.T) {
		initialRequests := bulk.pendingRequests
		docId := model.NewId()
		jsonData := []byte(`{"message": "test raw message"}`)

		err := bulk.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(docId),
		}, jsonData)
		require.NoError(t, err)
		require.Equal(t, initialRequests+1, bulk.pendingRequests)
	})

	t.Run("multiple index operations", func(t *testing.T) {
		initialRequests := bulk.pendingRequests

		for i := range 5 {
			post := createTestPost(t, fmt.Sprintf("test message %d", i))
			err := bulk.IndexOp(&types.IndexOperation{
				Index_: new("testindex"),
				Id_:    new(post.Id),
			}, post)
			require.NoError(t, err)
		}

		require.Equal(t, initialRequests+5, bulk.pendingRequests)
	})

	t.Run("auto flush on request threshold", func(t *testing.T) {
		// Create a new client with low flush threshold
		bulk2 := setupBulkClient(t, 0, 2, 0)
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		post1 := createTestPost(t, "first message")
		err := bulk2.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(post1.Id),
		}, post1)
		require.NoError(t, err)
		require.Equal(t, 1, bulk2.pendingRequests)

		post2 := createTestPost(t, "second message")
		err = bulk2.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(post2.Id),
		}, post2)
		require.NoError(t, err)
		require.Equal(t, 2, bulk2.pendingRequests)

		// Third operation should trigger flush
		post3 := createTestPost(t, "third message")
		err = bulk2.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(post3.Id),
		}, post3)
		require.NoError(t, err)
		require.Equal(t, 0, bulk2.pendingRequests)

		err = bulk2.Stop()
		require.NoError(t, err)
	})

	err := bulk.Stop()
	require.NoError(t, err)
}

func TestDeleteOp(t *testing.T) {
	bulk := setupBulkClient(t, 0, 10, 0)
	defer func() {
		if os.Getenv("IS_CI") == "true" {
			os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
			os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
		}
	}()

	t.Run("single delete operation", func(t *testing.T) {
		docId := model.NewId()

		err := bulk.DeleteOp(&types.DeleteOperation{
			Index_: new("testindex"),
			Id_:    new(docId),
		})
		require.NoError(t, err)
		require.Equal(t, 1, bulk.pendingRequests)

		// Verify buffer has content
		require.Greater(t, bulk.buf.Len(), 0)
	})

	t.Run("multiple delete operations", func(t *testing.T) {
		initialRequests := bulk.pendingRequests

		for range 3 {
			docId := model.NewId()
			err := bulk.DeleteOp(&types.DeleteOperation{
				Index_: new("testindex"),
				Id_:    new(docId),
			})
			require.NoError(t, err)
		}

		require.Equal(t, initialRequests+3, bulk.pendingRequests)
	})

	t.Run("auto flush on request threshold", func(t *testing.T) {
		// Create a new client with low flush threshold
		bulk2 := setupBulkClient(t, 0, 2, 0)
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		// Add two delete operations
		for range 2 {
			docId := model.NewId()
			err := bulk2.DeleteOp(&types.DeleteOperation{
				Index_: new("testindex"),
				Id_:    new(docId),
			})
			require.NoError(t, err)
		}
		require.Equal(t, 2, bulk2.pendingRequests)

		// Third operation should trigger flush
		docId := model.NewId()
		err := bulk2.DeleteOp(&types.DeleteOperation{
			Index_: new("testindex"),
			Id_:    new(docId),
		})
		require.NoError(t, err)
		require.Equal(t, 0, bulk2.pendingRequests)

		err = bulk2.Stop()
		require.NoError(t, err)
	})

	err := bulk.Stop()
	require.NoError(t, err)
}

func TestFlush(t *testing.T) {
	bulk := setupBulkClient(t, 0, 10, 0)
	defer func() {
		if os.Getenv("IS_CI") == "true" {
			os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
			os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
		}
	}()

	t.Run("flush with pending operations", func(t *testing.T) {
		post := createTestPost(t, "test message")

		err := bulk.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(post.Id),
		}, post)
		require.NoError(t, err)
		require.Equal(t, 1, bulk.pendingRequests)

		err = bulk.Flush()
		require.NoError(t, err)
		require.Equal(t, 0, bulk.pendingRequests)

		// Verify buffer is empty after flush
		require.Equal(t, 0, bulk.buf.Len())
	})

	t.Run("flush with no pending operations", func(t *testing.T) {
		require.Equal(t, 0, bulk.pendingRequests)

		err := bulk.Flush()
		require.NoError(t, err)
		require.Equal(t, 0, bulk.pendingRequests)
	})

	err := bulk.Stop()
	require.NoError(t, err)
}

func TestStop(t *testing.T) {
	t.Run("stop with pending operations", func(t *testing.T) {
		bulk := setupBulkClient(t, 0, 10, 0)
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		post := createTestPost(t, "test message")

		err := bulk.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(post.Id),
		}, post)
		require.NoError(t, err)
		require.Equal(t, 1, bulk.pendingRequests)

		err = bulk.Stop()
		require.NoError(t, err)
		require.Equal(t, 0, bulk.pendingRequests)
	})

	t.Run("stop with no pending operations", func(t *testing.T) {
		bulk := setupBulkClient(t, 0, 10, 0)
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		require.Equal(t, 0, bulk.pendingRequests)

		err := bulk.Stop()
		require.NoError(t, err)
		require.Equal(t, 0, bulk.pendingRequests)
	})

	t.Run("stop with periodic flusher", func(t *testing.T) {
		bulk := setupBulkClient(t, 0, 10, 100*time.Millisecond)
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		post := createTestPost(t, "test message")

		err := bulk.IndexOp(&types.IndexOperation{
			Index_: new("testindex"),
			Id_:    new(post.Id),
		}, post)
		require.NoError(t, err)
		require.Equal(t, 1, bulk.pendingRequests)

		// Stop should flush pending operations and stop the periodic flusher
		err = bulk.Stop()
		require.NoError(t, err)
		require.Equal(t, 0, bulk.pendingRequests)
	})
}

func TestStopShutsDownPeriodicFlusher(t *testing.T) {
	tests := []struct {
		name            string
		status          int
		wantErr         bool
		pendingRequests int
	}{
		{name: "successful pending flush", status: http.StatusOK},
		{name: "failed pending flush", status: http.StatusInternalServerError, wantErr: true, pendingRequests: 1},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			const flushInterval = 50 * time.Millisecond
			requestStarted := make(chan struct{}, 1)
			releaseRequest := make(chan struct{})
			var releaseOnce sync.Once
			release := func() {
				releaseOnce.Do(func() { close(releaseRequest) })
			}
			t.Cleanup(release)

			bulk := setupBulkClientWithHandler(t, flushInterval, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				select {
				case requestStarted <- struct{}{}:
				default:
				}
				<-releaseRequest

				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(tt.status)
				if tt.status == http.StatusOK {
					_, _ = w.Write([]byte(`{"errors":false,"items":[]}`))
					return
				}
				_, _ = w.Write([]byte(`{"error":{"type":"test_error","reason":"bulk request failed"},"status":500}`))
			}))

			post := createTestPost(t, "test message")
			err := bulk.IndexOp(&types.IndexOperation{
				Index_: new("testindex"),
				Id_:    new(post.Id),
			}, post)
			require.NoError(t, err)

			stopDone := make(chan error, 1)
			go func() {
				stopDone <- bulk.Stop()
			}()

			select {
			case <-requestStarted:
			case <-time.After(2 * time.Second):
				require.FailNow(t, "timed out waiting for final flush to start")
			}

			// Stop still holds mut while the final flush is blocked. Waiting for a
			// timer interval forces the periodic flusher to contend for that mutex.
			time.Sleep(2 * flushInterval)
			release()

			select {
			case err = <-stopDone:
			case <-time.After(2 * time.Second):
				require.FailNow(t, "Stop deadlocked while waiting for the periodic flusher")
			}

			if tt.wantErr {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
			}
			require.Equal(t, tt.pendingRequests, bulk.pendingRequests)
			assertFlusherStopped(t, bulk)
		})
	}
}

func setupBulkClientWithHandler(t *testing.T, flushInterval time.Duration, handler http.Handler) *Bulk {
	t.Helper()

	server := httptest.NewServer(handler)
	t.Cleanup(server.Close)
	client, err := opensearchapi.NewClient(opensearchapi.Config{
		Client: opensearch.Config{
			Addresses:    []string{server.URL},
			DisableRetry: true,
		},
	})
	require.NoError(t, err)

	return NewBulk(
		common.BulkSettings{
			FlushInterval: flushInterval,
			FlushNumReqs:  10,
		},
		client,
		time.Second,
		mlog.CreateConsoleTestLogger(t),
	)
}

func assertFlusherStopped(t *testing.T, bulk *Bulk) {
	t.Helper()

	select {
	case <-bulk.quitFlusher:
	default:
		require.Fail(t, "periodic flusher was not stopped")
	}
}

func TestFlushThresholds(t *testing.T) {
	t.Run("flush on bytes threshold", func(t *testing.T) {
		// Create a client with very small byte threshold
		bulk := setupBulkClient(t, 100, 0, 0) // 100 bytes threshold
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		// Add operations that should exceed the byte threshold
		for range 5 {
			post := createTestPost(t, "This is a long message that should help us exceed the byte threshold for testing")
			err := bulk.IndexOp(&types.IndexOperation{
				Index_: new("testindex"),
				Id_:    new(post.Id),
			}, post)
			require.NoError(t, err)
		}

		// Should have been flushed due to byte threshold
		require.Equal(t, 0, bulk.pendingRequests)

		err := bulk.Stop()
		require.NoError(t, err)
	})

	t.Run("no flush when thresholds not met", func(t *testing.T) {
		bulk := setupBulkClient(t, 100000, 10, 0) // High thresholds
		defer func() {
			if os.Getenv("IS_CI") == "true" {
				os.Setenv("MM_ELASTICSEARCHSETTINGS_CONNECTIONURL", "http://elasticsearch:9201")
				os.Unsetenv("MM_ELASTICSEARCHSETTINGS_BACKEND")
			}
		}()

		// Add a few operations that shouldn't trigger flush
		for range 3 {
			post := createTestPost(t, "short")
			err := bulk.IndexOp(&types.IndexOperation{
				Index_: new("testindex"),
				Id_:    new(post.Id),
			}, post)
			require.NoError(t, err)
			fmt.Println("PENDING REQS:", bulk.pendingRequests)
		}

		// Should not have been flushed
		require.Equal(t, 3, bulk.pendingRequests)

		err := bulk.Stop()
		require.NoError(t, err)
	})
}