github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/stores/indexshipper/compactor/table_test.go (about)

     1  package compactor
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io/ioutil"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  	"testing"
    11  
    12  	"github.com/go-kit/log"
    13  	"github.com/prometheus/common/model"
    14  	"github.com/stretchr/testify/require"
    15  
    16  	"github.com/grafana/loki/pkg/storage/chunk/client/local"
    17  	"github.com/grafana/loki/pkg/storage/config"
    18  	"github.com/grafana/loki/pkg/storage/stores/indexshipper/compactor/retention"
    19  	"github.com/grafana/loki/pkg/storage/stores/indexshipper/storage"
    20  )
    21  
    22  const (
    23  	objectsStorageDirName = "objects"
    24  	workingDirName        = "working-dir"
    25  	tableName             = "test"
    26  )
    27  
    28  type indexSetState struct {
    29  	uploadCompactedDB   bool
    30  	removeSourceObjects bool
    31  }
    32  
    33  func TestTable_Compaction(t *testing.T) {
    34  	for _, numUsers := range []int{uploadIndexSetsConcurrency / 2, uploadIndexSetsConcurrency, uploadIndexSetsConcurrency * 2} {
    35  		t.Run(fmt.Sprintf("numUsers=%d", numUsers), func(t *testing.T) {
    36  			for _, tc := range []struct {
    37  				numUnCompactedCommonDBs  int
    38  				numUnCompactedPerUserDBs int
    39  				numCompactedDBs          int
    40  
    41  				commonIndexSetState indexSetState
    42  				userIndexSetState   indexSetState
    43  			}{
    44  				{},
    45  				{
    46  					numCompactedDBs: 1,
    47  				},
    48  				{
    49  					numCompactedDBs: 2,
    50  					commonIndexSetState: indexSetState{
    51  						uploadCompactedDB:   true,
    52  						removeSourceObjects: true,
    53  					},
    54  					userIndexSetState: indexSetState{
    55  						uploadCompactedDB:   true,
    56  						removeSourceObjects: true,
    57  					},
    58  				},
    59  				{
    60  					numUnCompactedCommonDBs: 1,
    61  					commonIndexSetState: indexSetState{
    62  						uploadCompactedDB:   true,
    63  						removeSourceObjects: true,
    64  					},
    65  				},
    66  				{
    67  					numUnCompactedCommonDBs: 10,
    68  					commonIndexSetState: indexSetState{
    69  						uploadCompactedDB:   true,
    70  						removeSourceObjects: true,
    71  					},
    72  				},
    73  				{
    74  					numUnCompactedCommonDBs: 10,
    75  					numCompactedDBs:         1,
    76  					commonIndexSetState: indexSetState{
    77  						uploadCompactedDB:   true,
    78  						removeSourceObjects: true,
    79  					},
    80  				},
    81  				{
    82  					numUnCompactedCommonDBs: 10,
    83  					numCompactedDBs:         2,
    84  					commonIndexSetState: indexSetState{
    85  						uploadCompactedDB:   true,
    86  						removeSourceObjects: true,
    87  					},
    88  					userIndexSetState: indexSetState{
    89  						uploadCompactedDB:   true,
    90  						removeSourceObjects: true,
    91  					},
    92  				},
    93  				{
    94  					numUnCompactedPerUserDBs: 1,
    95  					commonIndexSetState: indexSetState{
    96  						removeSourceObjects: true,
    97  					},
    98  					userIndexSetState: indexSetState{
    99  						uploadCompactedDB:   true,
   100  						removeSourceObjects: true,
   101  					},
   102  				},
   103  				{
   104  					numUnCompactedPerUserDBs: 1,
   105  					numCompactedDBs:          1,
   106  					commonIndexSetState: indexSetState{
   107  						uploadCompactedDB:   true,
   108  						removeSourceObjects: true,
   109  					},
   110  					userIndexSetState: indexSetState{
   111  						uploadCompactedDB:   true,
   112  						removeSourceObjects: true,
   113  					},
   114  				},
   115  				{
   116  					numUnCompactedPerUserDBs: 1,
   117  					numCompactedDBs:          2,
   118  					commonIndexSetState: indexSetState{
   119  						uploadCompactedDB:   true,
   120  						removeSourceObjects: true,
   121  					},
   122  					userIndexSetState: indexSetState{
   123  						uploadCompactedDB:   true,
   124  						removeSourceObjects: true,
   125  					},
   126  				},
   127  				{
   128  					numUnCompactedPerUserDBs: 10,
   129  					commonIndexSetState: indexSetState{
   130  						removeSourceObjects: true,
   131  					},
   132  					userIndexSetState: indexSetState{
   133  						uploadCompactedDB:   true,
   134  						removeSourceObjects: true,
   135  					},
   136  				},
   137  				{
   138  					numUnCompactedCommonDBs:  10,
   139  					numUnCompactedPerUserDBs: 10,
   140  					commonIndexSetState: indexSetState{
   141  						uploadCompactedDB:   true,
   142  						removeSourceObjects: true,
   143  					},
   144  					userIndexSetState: indexSetState{
   145  						uploadCompactedDB:   true,
   146  						removeSourceObjects: true,
   147  					},
   148  				},
   149  				{
   150  					numUnCompactedCommonDBs:  10,
   151  					numUnCompactedPerUserDBs: 10,
   152  					numCompactedDBs:          1,
   153  					commonIndexSetState: indexSetState{
   154  						uploadCompactedDB:   true,
   155  						removeSourceObjects: true,
   156  					},
   157  					userIndexSetState: indexSetState{
   158  						uploadCompactedDB:   true,
   159  						removeSourceObjects: true,
   160  					},
   161  				},
   162  				{
   163  					numUnCompactedCommonDBs:  10,
   164  					numUnCompactedPerUserDBs: 10,
   165  					numCompactedDBs:          2,
   166  					commonIndexSetState: indexSetState{
   167  						uploadCompactedDB:   true,
   168  						removeSourceObjects: true,
   169  					},
   170  					userIndexSetState: indexSetState{
   171  						uploadCompactedDB:   true,
   172  						removeSourceObjects: true,
   173  					},
   174  				},
   175  			} {
   176  				commonDBsConfig := IndexesConfig{
   177  					NumCompactedFiles:   tc.numCompactedDBs,
   178  					NumUnCompactedFiles: tc.numUnCompactedCommonDBs,
   179  				}
   180  				perUserDBsConfig := PerUserIndexesConfig{
   181  					IndexesConfig: IndexesConfig{
   182  						NumCompactedFiles:   tc.numCompactedDBs,
   183  						NumUnCompactedFiles: tc.numUnCompactedPerUserDBs,
   184  					},
   185  					NumUsers: numUsers,
   186  				}
   187  
   188  				t.Run(fmt.Sprintf("%s ; %s", commonDBsConfig.String(), perUserDBsConfig.String()), func(t *testing.T) {
   189  					tempDir := t.TempDir()
   190  
   191  					objectStoragePath := filepath.Join(tempDir, objectsStorageDirName)
   192  					tablePathInStorage := filepath.Join(objectStoragePath, tableName)
   193  					tableWorkingDirectory := filepath.Join(tempDir, workingDirName, tableName)
   194  
   195  					SetupTable(t, filepath.Join(objectStoragePath, tableName), commonDBsConfig, perUserDBsConfig)
   196  
   197  					// do the compaction
   198  					objectClient, err := local.NewFSObjectClient(local.FSConfig{Directory: objectStoragePath})
   199  					require.NoError(t, err)
   200  
   201  					table, err := newTable(context.Background(), tableWorkingDirectory, storage.NewIndexStorageClient(objectClient, ""),
   202  						newTestIndexCompactor(), config.PeriodConfig{}, nil, nil)
   203  					require.NoError(t, err)
   204  
   205  					require.NoError(t, table.compact(false))
   206  
   207  					numUserIndexSets, numCommonIndexSets := 0, 0
   208  					for _, is := range table.indexSets {
   209  						if is.baseIndexSet.IsUserBasedIndexSet() {
   210  							require.Equal(t, tc.userIndexSetState.uploadCompactedDB, is.uploadCompactedDB)
   211  							require.Equal(t, tc.userIndexSetState.removeSourceObjects, is.removeSourceObjects)
   212  							numUserIndexSets++
   213  						} else {
   214  							require.Equal(t, tc.commonIndexSetState.uploadCompactedDB, is.uploadCompactedDB)
   215  							require.Equal(t, tc.commonIndexSetState.removeSourceObjects, is.removeSourceObjects)
   216  							numCommonIndexSets++
   217  						}
   218  					}
   219  
   220  					// verify the state in the storage after compaction.
   221  					expectedNumCommonDBs := 0
   222  					if (commonDBsConfig.NumUnCompactedFiles + commonDBsConfig.NumCompactedFiles) > 0 {
   223  						require.Equal(t, 1, numCommonIndexSets)
   224  						expectedNumCommonDBs = 1
   225  					}
   226  					numExpectedUsers := 0
   227  					if (perUserDBsConfig.NumUnCompactedFiles + perUserDBsConfig.NumCompactedFiles) > 0 {
   228  						require.Equal(t, numUsers, numUserIndexSets)
   229  						numExpectedUsers = numUsers
   230  					}
   231  					validateTable(t, tablePathInStorage, expectedNumCommonDBs, numExpectedUsers, func(filename string) {
   232  						require.True(t, strings.HasSuffix(filename, ".gz"), filename)
   233  					})
   234  
   235  					verifyCompactedIndexTable(t, commonDBsConfig, perUserDBsConfig, tablePathInStorage)
   236  
   237  					// running compaction again should not do anything.
   238  					table, err = newTable(context.Background(), tableWorkingDirectory, storage.NewIndexStorageClient(objectClient, ""),
   239  						newTestIndexCompactor(), config.PeriodConfig{}, nil, nil)
   240  					require.NoError(t, err)
   241  
   242  					require.NoError(t, table.compact(false))
   243  
   244  					for _, is := range table.indexSets {
   245  						require.False(t, is.uploadCompactedDB)
   246  						require.False(t, is.removeSourceObjects)
   247  					}
   248  				})
   249  			}
   250  		})
   251  	}
   252  }
   253  
   254  type TableMarkerFunc func(ctx context.Context, tableName, userID string, indexFile retention.IndexProcessor, logger log.Logger) (bool, bool, error)
   255  
   256  func (t TableMarkerFunc) MarkForDelete(ctx context.Context, tableName, userID string, indexFile retention.IndexProcessor, logger log.Logger) (bool, bool, error) {
   257  	return t(ctx, tableName, userID, indexFile, logger)
   258  }
   259  
   260  type IntervalMayHaveExpiredChunksFunc func(interval model.Interval, userID string) bool
   261  
   262  func (f IntervalMayHaveExpiredChunksFunc) IntervalMayHaveExpiredChunks(interval model.Interval, userID string) bool {
   263  	return f(interval, userID)
   264  }
   265  
   266  func TestTable_CompactionRetention(t *testing.T) {
   267  	numUsers := 10
   268  	type dbsSetup struct {
   269  		numUnCompactedCommonDBs  int
   270  		numUnCompactedPerUserDBs int
   271  		numCompactedDBs          int
   272  	}
   273  	for _, setup := range []dbsSetup{
   274  		{
   275  			numUnCompactedCommonDBs:  10,
   276  			numUnCompactedPerUserDBs: 10,
   277  		},
   278  		{
   279  			numCompactedDBs: 1,
   280  		},
   281  		{
   282  			numCompactedDBs: 10,
   283  		},
   284  		{
   285  			numUnCompactedCommonDBs:  10,
   286  			numUnCompactedPerUserDBs: 10,
   287  			numCompactedDBs:          1,
   288  		},
   289  		{
   290  			numUnCompactedCommonDBs: 1,
   291  		},
   292  		{
   293  			numUnCompactedPerUserDBs: 1,
   294  		},
   295  	} {
   296  		for name, tt := range map[string]struct {
   297  			dbsSetup    dbsSetup
   298  			dbCount     int
   299  			assert      func(t *testing.T, storagePath, tableName string)
   300  			tableMarker retention.TableMarker
   301  		}{
   302  			"emptied table": {
   303  				dbsSetup: setup,
   304  				assert: func(t *testing.T, storagePath, tableName string) {
   305  					_, err := ioutil.ReadDir(filepath.Join(storagePath, tableName))
   306  					require.True(t, os.IsNotExist(err))
   307  				},
   308  				tableMarker: TableMarkerFunc(func(ctx context.Context, tableName, userID string, indexFile retention.IndexProcessor, logger log.Logger) (bool, bool, error) {
   309  					return true, true, nil
   310  				}),
   311  			},
   312  			"marked table": {
   313  				dbsSetup: setup,
   314  				assert: func(t *testing.T, storagePath, tableName string) {
   315  					expectedNumCommonDBs := 0
   316  					if setup.numUnCompactedCommonDBs+setup.numCompactedDBs > 0 {
   317  						expectedNumCommonDBs = 1
   318  					}
   319  
   320  					expectedNumUsers := 0
   321  					if setup.numUnCompactedPerUserDBs+setup.numCompactedDBs > 0 {
   322  						expectedNumUsers = numUsers
   323  					}
   324  					validateTable(t, filepath.Join(storagePath, tableName), expectedNumCommonDBs, expectedNumUsers, func(filename string) {
   325  						require.True(t, strings.HasSuffix(filename, ".gz"))
   326  					})
   327  				},
   328  				tableMarker: TableMarkerFunc(func(ctx context.Context, tableName, userID string, indexFile retention.IndexProcessor, logger log.Logger) (bool, bool, error) {
   329  					return false, true, nil
   330  				}),
   331  			},
   332  			"not modified": {
   333  				dbsSetup: setup,
   334  				assert: func(t *testing.T, storagePath, tableName string) {
   335  					expectedNumCommonDBs := 0
   336  					if setup.numUnCompactedCommonDBs+setup.numCompactedDBs > 0 {
   337  						expectedNumCommonDBs = 1
   338  					}
   339  
   340  					expectedNumUsers := 0
   341  					if setup.numUnCompactedPerUserDBs+setup.numCompactedDBs > 0 {
   342  						expectedNumUsers = numUsers
   343  					}
   344  					validateTable(t, filepath.Join(storagePath, tableName), expectedNumCommonDBs, expectedNumUsers, func(filename string) {
   345  						require.True(t, strings.HasSuffix(filename, ".gz"))
   346  					})
   347  				},
   348  				tableMarker: TableMarkerFunc(func(ctx context.Context, tableName, userID string, indexFile retention.IndexProcessor, logger log.Logger) (bool, bool, error) {
   349  					return false, false, nil
   350  				}),
   351  			},
   352  		} {
   353  			tt := tt
   354  			commonDBsConfig := IndexesConfig{
   355  				NumCompactedFiles:   tt.dbsSetup.numCompactedDBs,
   356  				NumUnCompactedFiles: tt.dbsSetup.numUnCompactedCommonDBs,
   357  			}
   358  			perUserDBsConfig := PerUserIndexesConfig{
   359  				IndexesConfig: IndexesConfig{
   360  					NumUnCompactedFiles: tt.dbsSetup.numUnCompactedPerUserDBs,
   361  					NumCompactedFiles:   tt.dbsSetup.numCompactedDBs,
   362  				},
   363  				NumUsers: numUsers,
   364  			}
   365  			t.Run(fmt.Sprintf("%s - %s ; %s", name, commonDBsConfig.String(), perUserDBsConfig.String()), func(t *testing.T) {
   366  				tempDir := t.TempDir()
   367  				tableName := fmt.Sprintf("%s12345", tableName)
   368  
   369  				objectStoragePath := filepath.Join(tempDir, objectsStorageDirName)
   370  				tableWorkingDirectory := filepath.Join(tempDir, workingDirName, tableName)
   371  
   372  				SetupTable(t, filepath.Join(objectStoragePath, tableName), commonDBsConfig, perUserDBsConfig)
   373  
   374  				// do the compaction
   375  				objectClient, err := local.NewFSObjectClient(local.FSConfig{Directory: objectStoragePath})
   376  				require.NoError(t, err)
   377  
   378  				table, err := newTable(context.Background(), tableWorkingDirectory, storage.NewIndexStorageClient(objectClient, ""),
   379  					newTestIndexCompactor(), config.PeriodConfig{},
   380  					tt.tableMarker, IntervalMayHaveExpiredChunksFunc(func(interval model.Interval, userID string) bool {
   381  						return true
   382  					}))
   383  				require.NoError(t, err)
   384  
   385  				require.NoError(t, table.compact(true))
   386  				tt.assert(t, objectStoragePath, tableName)
   387  			})
   388  		}
   389  	}
   390  }
   391  
   392  func validateTable(t *testing.T, path string, expectedNumCommonDBs, numUsers int, filesCallback func(filename string)) {
   393  	files, folders := listDir(t, path)
   394  	require.Len(t, files, expectedNumCommonDBs)
   395  	require.Len(t, folders, numUsers)
   396  
   397  	for _, fileName := range files {
   398  		filesCallback(fileName)
   399  	}
   400  
   401  	for _, folder := range folders {
   402  		files, folders := listDir(t, filepath.Join(path, folder))
   403  		require.Len(t, files, 1)
   404  		require.Len(t, folders, 0)
   405  
   406  		for _, fileName := range files {
   407  			filesCallback(fileName)
   408  		}
   409  	}
   410  }
   411  
   412  func listDir(t *testing.T, path string) (files, folders []string) {
   413  	filesInfo, err := ioutil.ReadDir(path)
   414  	require.NoError(t, err)
   415  
   416  	for _, fileInfo := range filesInfo {
   417  		if fileInfo.IsDir() {
   418  			folders = append(folders, fileInfo.Name())
   419  		} else {
   420  			files = append(files, fileInfo.Name())
   421  		}
   422  	}
   423  
   424  	return
   425  }
   426  
   427  func TestTable_CompactionFailure(t *testing.T) {
   428  	tempDir := t.TempDir()
   429  
   430  	tableName := "test"
   431  	objectStoragePath := filepath.Join(tempDir, objectsStorageDirName)
   432  	tablePathInStorage := filepath.Join(objectStoragePath, tableName)
   433  	tableWorkingDirectory := filepath.Join(tempDir, workingDirName, tableName)
   434  
   435  	// setup some dbs
   436  	numDBs := 10
   437  
   438  	dbsToSetup := make(map[string]IndexFileConfig)
   439  	for i := 0; i < numDBs; i++ {
   440  		dbsToSetup[fmt.Sprint(i)] = IndexFileConfig{
   441  			CompressFile: i%2 == 0,
   442  		}
   443  	}
   444  
   445  	SetupTable(t, filepath.Join(objectStoragePath, tableName), IndexesConfig{NumCompactedFiles: numDBs}, PerUserIndexesConfig{})
   446  
   447  	// put a corrupt zip file in the table which should cause the compaction to fail in the middle because it would fail to open that file with boltdb client.
   448  	require.NoError(t, ioutil.WriteFile(filepath.Join(tablePathInStorage, "fail.gz"), []byte("fail the compaction"), 0o666))
   449  
   450  	// do the compaction
   451  	objectClient, err := local.NewFSObjectClient(local.FSConfig{Directory: objectStoragePath})
   452  	require.NoError(t, err)
   453  
   454  	table, err := newTable(context.Background(), tableWorkingDirectory, storage.NewIndexStorageClient(objectClient, ""),
   455  		newTestIndexCompactor(), config.PeriodConfig{}, nil, nil)
   456  	require.NoError(t, err)
   457  
   458  	// compaction should fail due to a non-boltdb file.
   459  	require.Error(t, table.compact(false))
   460  
   461  	// ensure that files in storage are intact.
   462  	files, err := ioutil.ReadDir(tablePathInStorage)
   463  	require.NoError(t, err)
   464  	require.Len(t, files, numDBs+1)
   465  
   466  	// ensure that we have cleanup the local working directory after failing the compaction.
   467  	require.NoFileExists(t, tableWorkingDirectory)
   468  
   469  	// remove the corrupt zip file and ensure that compaction succeeds now.
   470  	require.NoError(t, os.Remove(filepath.Join(tablePathInStorage, "fail.gz")))
   471  
   472  	table, err = newTable(context.Background(), tableWorkingDirectory, storage.NewIndexStorageClient(objectClient, ""),
   473  		newTestIndexCompactor(), config.PeriodConfig{}, nil, nil)
   474  	require.NoError(t, err)
   475  	require.NoError(t, table.compact(false))
   476  
   477  	// ensure that we have cleanup the local working directory after successful compaction.
   478  	require.NoFileExists(t, tableWorkingDirectory)
   479  }