github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/tools/blocksconvert/scheduler/scheduler_test.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"regexp"
     8  	"strings"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/go-kit/log"
    13  	"github.com/stretchr/testify/require"
    14  	"github.com/thanos-io/thanos/pkg/objstore"
    15  
    16  	"github.com/cortexproject/cortex/tools/blocksconvert"
    17  )
    18  
    19  func TestScanForPlans(t *testing.T) {
    20  	bucket := objstore.NewInMemBucket()
    21  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/1.plan", strings.NewReader("")))
    22  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/1.starting.1234567", strings.NewReader("")))
    23  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/1.inprogress.2345678", strings.NewReader("")))
    24  
    25  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/2.plan", strings.NewReader("")))
    26  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/2.inprogress.93485345", strings.NewReader("")))
    27  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/2.finished.01E8GCW9J0HV0992HSZ0N6RAMN", strings.NewReader("")))
    28  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/2.finished.01EE9Y140JP4T58X8FGTG5T17F", strings.NewReader("")))
    29  
    30  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/3.plan", strings.NewReader("")))
    31  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/3.error", strings.NewReader("")))
    32  
    33  	// Only error, progress or finished
    34  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/4.error", strings.NewReader("")))
    35  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/5.progress.1234234", strings.NewReader("")))
    36  	require.NoError(t, bucket.Upload(context.Background(), "migration/12345/6.finished.cleaned", strings.NewReader("")))
    37  
    38  	plans, err := scanForPlans(context.Background(), bucket, "migration/12345/")
    39  	require.NoError(t, err)
    40  
    41  	require.Equal(t, map[string]plan{
    42  		"1": {
    43  			PlanFiles: []string{"migration/12345/1.plan"},
    44  			ProgressFiles: map[string]time.Time{
    45  				"migration/12345/1.starting.1234567":   time.Unix(1234567, 0),
    46  				"migration/12345/1.inprogress.2345678": time.Unix(2345678, 0),
    47  			},
    48  		},
    49  		"2": {
    50  			PlanFiles: []string{"migration/12345/2.plan"},
    51  			ProgressFiles: map[string]time.Time{
    52  				"migration/12345/2.inprogress.93485345": time.Unix(93485345, 0),
    53  			},
    54  			Finished: []string{"01E8GCW9J0HV0992HSZ0N6RAMN", "01EE9Y140JP4T58X8FGTG5T17F"},
    55  		},
    56  		"3": {
    57  			PlanFiles: []string{"migration/12345/3.plan"},
    58  			ErrorFile: "migration/12345/3.error",
    59  		},
    60  		"4": {
    61  			ErrorFile: "migration/12345/4.error",
    62  		},
    63  		"5": {
    64  			ProgressFiles: map[string]time.Time{
    65  				"migration/12345/5.progress.1234234": time.Unix(1234234, 0),
    66  			},
    67  		},
    68  		"6": {
    69  			Finished: []string{"cleaned"},
    70  		},
    71  	}, plans)
    72  }
    73  
    74  func TestSchedulerScan(t *testing.T) {
    75  	now := time.Now()
    76  	nowMinus1Hour := now.Add(-time.Hour)
    77  
    78  	bucket := objstore.NewInMemBucket()
    79  	require.NoError(t, bucket.Upload(context.Background(), "migration/user1/1.plan", strings.NewReader("")))
    80  	// This progress file is too old, will be removed.
    81  	require.NoError(t, bucket.Upload(context.Background(), fmt.Sprintf("migration/user1/1.inprogress.%d", nowMinus1Hour.Unix()), strings.NewReader("")))
    82  
    83  	require.NoError(t, bucket.Upload(context.Background(), "migration/user2/2.plan", strings.NewReader("")))
    84  	require.NoError(t, bucket.Upload(context.Background(), fmt.Sprintf("migration/user2/2.inprogress.%d", now.Unix()), strings.NewReader("")))
    85  
    86  	require.NoError(t, bucket.Upload(context.Background(), "migration/user3/3.plan", strings.NewReader("")))
    87  	require.NoError(t, bucket.Upload(context.Background(), "migration/user3/3.error", strings.NewReader("")))
    88  
    89  	require.NoError(t, bucket.Upload(context.Background(), "migration/user4/4.plan", strings.NewReader("")))
    90  	require.NoError(t, bucket.Upload(context.Background(), "migration/user4/5.error", strings.NewReader("")))
    91  	require.NoError(t, bucket.Upload(context.Background(), "migration/user4/6.finished.01E8GCW9J0HV0992HSZ0N6RAMN", strings.NewReader("")))
    92  
    93  	require.NoError(t, bucket.Upload(context.Background(), "migration/ignoredUser/7.plan", strings.NewReader("")))
    94  
    95  	ignoredUsers := regexp.MustCompile("ignored.*")
    96  
    97  	s := newSchedulerWithBucket(log.NewLogfmtLogger(os.Stdout), bucket, "migration", blocksconvert.AllowAllUsers, ignoredUsers, Config{
    98  		ScanInterval:        10 * time.Second,
    99  		PlanScanConcurrency: 5,
   100  		MaxProgressFileAge:  5 * time.Minute,
   101  	}, nil)
   102  
   103  	require.NoError(t, s.scanBucketForPlans(context.Background()))
   104  	require.Equal(t, []queuedPlan{
   105  		{DayIndex: 4, PlanFile: "migration/user4/4.plan"},
   106  	}, s.plansQueue)
   107  	require.Equal(t, "migration/user1/1.error", s.allUserPlans["user1"]["1"].ErrorFile)
   108  
   109  	{
   110  		p, pg := s.nextPlanNoRunningCheck(context.Background())
   111  		require.Equal(t, "migration/user4/4.plan", p)
   112  		ok, err := bucket.Exists(context.Background(), pg)
   113  		require.NoError(t, err)
   114  		require.True(t, ok)
   115  	}
   116  
   117  	{
   118  		p, pg := s.nextPlanNoRunningCheck(context.Background())
   119  		require.Equal(t, "", p)
   120  		require.Equal(t, "", pg)
   121  	}
   122  }