<commit_before>package jobs

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"sync/atomic"
	"time"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/prefixer"
	"github.com/go-redis/redis"
	multierror "github.com/hashicorp/go-multierror"
)

const (
	// redisPrefix is the prefix for jobs queues in redis.
	redisPrefix = "j/"
	// redisHighPrioritySuffix suffix is the suffix used for prioritized queue.
	redisHighPrioritySuffix = "/p0"
)

type redisBroker struct {
	client       redis.UniversalClient
	workers      []*Worker
	workersTypes []string
	running      uint32
	closed       chan struct{}
}

// NewRedisBroker creates a new broker that will use redis to distribute
// the jobs among several cozy-stack processes.
func NewRedisBroker(client redis.UniversalClient) Broker {
	return &redisBroker{
		client: client,
		closed: make(chan struct{}),
	}
}

// StartWorkers polling jobs from redis queues
func (b *redisBroker) StartWorkers(ws WorkersList) error {
	if !atomic.CompareAndSwapUint32(&b.running, 0, 1) {
		return ErrClosed
	}

	for _, conf := range ws {
		b.workersTypes = append(b.workersTypes, conf.WorkerType)
		if conf.Concurrency <= 0 {
			continue
		}
		ch := make(chan *Job)
		w := NewWorker(conf)
		b.workers = append(b.workers, w)
		if err := w.Start(ch); err != nil {
			return err
		}
		go b.pollLoop(redisPrefix+conf.WorkerType, ch)
	}

	if len(b.workers) > 0 {
		joblog.Infof("Started redis broker for %d workers type", len(b.workers))
	}

	// XXX for retro-compat
	if slots := config.GetConfig().Jobs.NbWorkers; len(b.workers) > 0 && slots > 0 {
		joblog.Warnf("Limiting the number of total concurrent workers to %d", slots)
		joblog.Warnf("Please update your configuration file to avoid a hard limit")
		setNbSlots(slots)
	}

	return nil
}

func (b *redisBroker) WorkersTypes() []string {
	return b.workersTypes
}

func (b *redisBroker) ShutdownWorkers(ctx context.Context) error {
	if !atomic.CompareAndSwapUint32(&b.running, 1, 0) {
		return ErrClosed
	}
	if len(b.workers) == 0 {
		return nil
	}

	fmt.Print("  shutting down redis broker...")
	defer b.client.Close()

	for i := 0; i < len(b.workers); i++ {
		select {
		case <-ctx.Done():
			fmt.Println("failed:", ctx.Err())
			return ctx.Err()
		case <-b.closed:
		}
	}

	errs := make(chan error)
	for _, w := range b.workers {
		go func(w *Worker) { errs <- w.Shutdown(ctx) }(w)
	}

	var errm error
	for i := 0; i < len(b.workers); i++ {
		if err := <-errs; err != nil {
			errm = multierror.Append(errm, err)
		}
	}

	if errm != nil {
		fmt.Println("failed: ", errm)
	} else {
		fmt.Println("ok")
	}
	return errm
}

var redisBRPopTimeout = 10 * time.Second

func (b *redisBroker) pollLoop(key string, ch chan<- *Job) {
	defer func() {
		b.closed <- struct{}{}<commit_msg>Fix: allow starting jobs from workers that are not running locally<commit_after>package jobs

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"sync/atomic"
	"time"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/prefixer"
	"github.com/go-redis/redis"
	multierror "github.com/hashicorp/go-multierror"
)

const (
	// redisPrefix is the prefix for jobs queues in redis.
	redisPrefix = "j/"
	// redisHighPrioritySuffix suffix is the suffix used for prioritized queue.
	redisHighPrioritySuffix = "/p0"
)

type redisBroker struct {
	client         redis.UniversalClient
	workers        []*Worker
	workersRunning []*Worker
	workersTypes   []string
	running        uint32
	closed         chan struct{}
}

// NewRedisBroker creates a new broker that will use redis to distribute
// the jobs among several cozy-stack processes.
func NewRedisBroker(client redis.UniversalClient) Broker {
	return &redisBroker{
		client: client,
		closed: make(chan struct{}),
	}
}

// StartWorkers polling jobs from redis queues
func (b *redisBroker) StartWorkers(ws WorkersList) error {
	if !atomic.CompareAndSwapUint32(&b.running, 0, 1) {
		return ErrClosed
	}

	for _, conf := range ws {
		b.workersTypes = append(b.workersTypes, conf.WorkerType)
		w := NewWorker(conf)
		b.workers = append(b.workers, w)
		if conf.Concurrency <= 0 {
			continue
		}
		b.workersRunning = append(b.workersRunning, w)
		ch := make(chan *Job)
		if err := w.Start(ch); err != nil {
			return err
		}
		go b.pollLoop(redisPrefix+conf.WorkerType, ch)
	}

	if len(b.workersRunning) > 0 {
		joblog.Infof("Started redis broker for %d workers type", len(b.workersRunning))
	}

	// XXX for retro-compat
	if slots := config.GetConfig().Jobs.NbWorkers; len(b.workersRunning) > 0 && slots > 0 {
		joblog.Warnf("Limiting the number of total concurrent workers to %d", slots)
		joblog.Warnf("Please update your configuration file to avoid a hard limit")
		setNbSlots(slots)
	}

	return nil
}

func (b *redisBroker) WorkersTypes() []string {
	return b.workersTypes
}

func (b *redisBroker) ShutdownWorkers(ctx context.Context) error {
	if !atomic.CompareAndSwapUint32(&b.running, 1, 0) {
		return ErrClosed
	}
	if len(b.workersRunning) == 0 {
		return nil
	}

	fmt.Print("  shutting down redis broker...")
	defer b.client.Close()

	for i := 0; i < len(b.workersRunning); i++ {
		select {
		case <-ctx.Done():
			fmt.Println("failed:", ctx.Err())
			return ctx.Err()
		case <-b.closed:
		}
	}

	errs := make(chan error)
	for _, w := range b.workersRunning {
		go func(w *Worker) { errs <- w.Shutdown(ctx) }(w)
	}

	var errm error
	for i := 0; i < len(b.workersRunning); i++ {
		if err := <-errs; err != nil {
			errm = multierror.Append(errm, err)
		}
	}

	if errm != nil {
		fmt.Println("failed: ", errm)
	} else {
		fmt.Println("ok")
	}
	return errm
}

var redisBRPopTimeout = 10 * time.Second

func (b *redisBroker) pollLoop(key string, ch chan<- *Job) {
	defer func() {
		b.closed <- struct{}{}
{'fraction_non_alphanumeric': 0.1356285863328117, 'fraction_numerical': 0.006955312119631368, 'mean_word_length': 3.5578446909667196, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3709614', 'n_tokens_mistral': 2258, 'n_tokens_neox': 2091, 'n_words': 556}