github.com/cozy/cozy-stack@v0.0.0-20240603063001-31110fa4cae1/model/job/redis_broker_test.go (about)

     1  package job_test
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"strconv"
     7  	"strings"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/cozy/cozy-stack/model/instance"
    13  	"github.com/cozy/cozy-stack/model/job"
    14  	"github.com/cozy/cozy-stack/pkg/config/config"
    15  	"github.com/cozy/cozy-stack/pkg/limits"
    16  	"github.com/cozy/cozy-stack/tests/testutils"
    17  	"github.com/redis/go-redis/v9"
    18  	"github.com/stretchr/testify/assert"
    19  )
    20  
    21  const redisURL1 = "redis://localhost:6379/0"
    22  const redisURL2 = "redis://localhost:6379/1"
    23  
    24  func TestRedisBroker(t *testing.T) {
    25  	if testing.Short() {
    26  		t.Skip("an instance is required for this test: test skipped due to the use of --short flag")
    27  	}
    28  
    29  	config.UseTestFile(t)
    30  	setup := testutils.NewSetup(t, t.Name())
    31  	testInstance := setup.GetTestInstance()
    32  
    33  	t.Run("RedisJobs", func(t *testing.T) {
    34  		job.SetRedisTimeoutForTest()
    35  		opts1, _ := redis.ParseURL(redisURL1)
    36  		opts2, _ := redis.ParseURL(redisURL2)
    37  		client1 := redis.NewClient(opts1)
    38  		client2 := redis.NewClient(opts2)
    39  
    40  		n := 10
    41  		v := 100
    42  
    43  		var w sync.WaitGroup
    44  		w.Add(2*n + 1)
    45  
    46  		workersTestList := job.WorkersList{
    47  			{
    48  				WorkerType:  "test",
    49  				Concurrency: 4,
    50  				WorkerFunc: func(ctx *job.TaskContext) error {
    51  					var msg string
    52  					err := ctx.UnmarshalMessage(&msg)
    53  					if !assert.NoError(t, err) {
    54  						return err
    55  					}
    56  					if strings.HasPrefix(msg, "z-") {
    57  						_, err := strconv.Atoi(msg[len("z-"):])
    58  						assert.NoError(t, err)
    59  					} else if strings.HasPrefix(msg, "a-") {
    60  						_, err := strconv.Atoi(msg[len("a-"):])
    61  						assert.NoError(t, err)
    62  					} else if strings.HasPrefix(msg, "b-") {
    63  						_, err := strconv.Atoi(msg[len("b-"):])
    64  						assert.NoError(t, err)
    65  					} else {
    66  						t.Fatal()
    67  					}
    68  					w.Done()
    69  					return nil
    70  				},
    71  			},
    72  		}
    73  
    74  		broker1 := job.NewRedisBroker(client1)
    75  		err := broker1.StartWorkers(workersTestList)
    76  		assert.NoError(t, err)
    77  
    78  		broker2 := job.NewRedisBroker(client2)
    79  		err = broker2.StartWorkers(workersTestList)
    80  		assert.NoError(t, err)
    81  
    82  		msg, _ := job.NewMessage("z-0")
    83  		_, err = broker1.PushJob(testInstance, &job.JobRequest{
    84  			WorkerType: "test",
    85  			Message:    msg,
    86  		})
    87  		assert.NoError(t, err)
    88  
    89  		go func(broker job.Broker, instance *instance.Instance, n int) {
    90  			for i := 0; i < n; i++ {
    91  				msg, _ := job.NewMessage("a-" + strconv.Itoa(i+1))
    92  				_, err2 := broker.PushJob(instance, &job.JobRequest{
    93  					WorkerType: "test",
    94  					Message:    msg,
    95  				})
    96  				assert.NoError(t, err2)
    97  				time.Sleep(randomMicro(0, v))
    98  			}
    99  		}(broker1, testInstance, n)
   100  
   101  		go func(broker job.Broker, instance *instance.Instance, n int) {
   102  			for i := 0; i < n; i++ {
   103  				msg, _ := job.NewMessage("b-" + strconv.Itoa(i+1))
   104  				_, err2 := broker.PushJob(instance, &job.JobRequest{
   105  					WorkerType: "test",
   106  					Message:    msg,
   107  					Manual:     true,
   108  				})
   109  				assert.NoError(t, err2)
   110  				time.Sleep(randomMicro(0, v))
   111  			}
   112  		}(broker2, testInstance, n)
   113  
   114  		w.Wait()
   115  
   116  		err = broker1.ShutdownWorkers(context.Background())
   117  		assert.NoError(t, err)
   118  		err = broker2.ShutdownWorkers(context.Background())
   119  		assert.NoError(t, err)
   120  		time.Sleep(1 * time.Second)
   121  	})
   122  
   123  	t.Run("RedisAddJobRateLimitExceeded", func(t *testing.T) {
   124  		opts1, _ := redis.ParseURL(redisURL1)
   125  		client1 := redis.NewClient(opts1)
   126  		workersTestList := job.WorkersList{
   127  			{
   128  				WorkerType:  "thumbnail",
   129  				Concurrency: 4,
   130  				WorkerFunc: func(ctx *job.TaskContext) error {
   131  					return nil
   132  				},
   133  			},
   134  		}
   135  		ct := limits.JobThumbnailType
   136  		config.GetRateLimiter().ResetCounter(testInstance, ct)
   137  
   138  		broker := job.NewRedisBroker(client1)
   139  		err := broker.StartWorkers(workersTestList)
   140  		assert.NoError(t, err)
   141  
   142  		msg, _ := job.NewMessage("z-0")
   143  		j, err := broker.PushJob(testInstance, &job.JobRequest{
   144  			WorkerType: "thumbnail",
   145  			Message:    msg,
   146  		})
   147  
   148  		assert.NoError(t, err)
   149  		assert.NotNil(t, j)
   150  
   151  		limits.SetMaximumLimit(ct, 10)
   152  		maxLimit := limits.GetMaximumLimit(ct)
   153  
   154  		// Blocking the job push
   155  		for i := int64(0); i < maxLimit-1; i++ {
   156  			j, err := broker.PushJob(testInstance, &job.JobRequest{
   157  				WorkerType: "thumbnail",
   158  				Message:    msg,
   159  			})
   160  			assert.NoError(t, err)
   161  			assert.NotNil(t, j)
   162  		}
   163  
   164  		j, err = broker.PushJob(testInstance, &job.JobRequest{
   165  			WorkerType: "thumbnail",
   166  			Message:    msg,
   167  		})
   168  		assert.Error(t, err)
   169  		assert.Nil(t, j)
   170  		assert.ErrorIs(t, err, limits.ErrRateLimitReached)
   171  
   172  		j, err = broker.PushJob(testInstance, &job.JobRequest{
   173  			WorkerType: "thumbnail",
   174  			Message:    msg,
   175  		})
   176  		assert.Error(t, err)
   177  		assert.Nil(t, j)
   178  	})
   179  }
   180  
   181  func randomMicro(min, max int) time.Duration {
   182  	return time.Duration(rand.Intn(max-min)+min) * time.Microsecond
   183  }