<commit_before>package brokers

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/signatures"

	"github.com/garyburd/redigo/redis"

	"strings"
)

// RedisBroker represents a Redis broker
type RedisBroker struct {
	config   *config.Config
	stopChan chan int
}

// NewRedisBroker creates new RedisBroker instance
func NewRedisBroker(cnf *config.Config) Broker {
	return Broker(&RedisBroker{
		config: cnf,
	})
}

// StartConsuming enters a loop and waits for incoming messages
func (redisBroker *RedisBroker) StartConsuming(consumerTag string, taskProcessor TaskProcessor) (bool, error) {
	conn, err := openRedisConn(redisBroker.config)
	if err != nil {
		return true, err // retry true
	}

	defer conn.Close()

	redisBroker.stopChan = make(chan int)

	psc := redis.PubSubConn{Conn: conn}
	if err := psc.Subscribe(redisBroker.config.DefaultQueue); err != nil {
		return true, err // retry true
	}
	// Unsubscribe from all connections. This will cause the receiving
	// goroutine to exit.
	defer psc.Unsubscribe()

	deliveries := make(chan signatures.TaskSignature)
	errors := make(chan error)

	log.Print("[*] Waiting for messages. To exit press CTRL+C")

	go func() {
		for {
			switch n := psc.Receive().(type) {
			case redis.Message:
				log.Printf("Received new message: %s", n.Data)

				signature := signatures.TaskSignature{}
				if err := json.Unmarshal(n.Data, &signature); err != nil {
					errors <- err
					continue
				}

				deliveries <- signature
			case redis.Subscription:
				if n.Count == 0 {
					return
				}
			case error:
				log.Print(n)
			}
		}
	}()

	for {
		select {
		case signature := <-deliveries:
			taskProcessor.Process(&signature)
		case err := <-errors:
			// Unsubscribe from all connections. This will cause the receiving
			// goroutine to exit.
			psc.Unsubscribe()
			return true, err
		case <-redisBroker.stopChan:
			// Unsubscribe from all connections. This will cause the receiving
			// goroutine to exit.
			psc.Unsubscribe()
			return false, nil
		}
	}
}

// StopConsuming quits the loop
func (redisBroker *RedisBroker) StopConsuming() {
	// Notifying the quit channel stops consuming of messages
	redisBroker.stopChan <- 1
}

// Publish places a new message on the default queue
func (redisBroker *RedisBroker) Publish(signature *signatures.TaskSignature) error {
	conn, err := openRedisConn(redisBroker.config)
	if err != nil {
		return err
	}

	defer conn.Close()

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON Encode Message: %v", err)
	}

	conn.Do("PUBLISH", redisBroker.config.DefaultQueue, message)
	return conn.Flush()
}

func openRedisConn(cnf *config.Config) (redis.Conn, error) {
	network := "tcp"
	address := strings.Split(cnf.Broker, "redis://")[1]
	return redis.Dial(network, address)
}<commit_msg>Refactoring Redis broker<commit_after>package brokers

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/signatures"

	"github.com/garyburd/redigo/redis"

	"strings"
)

// RedisBroker represents a Redis broker
type RedisBroker struct {
	config   *config.Config
	stopChan chan int
}

// NewRedisBroker creates new RedisBroker instance
func NewRedisBroker(cnf *config.Config) Broker {
	return Broker(&RedisBroker{
		config: cnf,
	})
}

// StartConsuming enters a loop and waits for incoming messages
func (redisBroker *RedisBroker) StartConsuming(consumerTag string, taskProcessor TaskProcessor) (bool, error) {
	conn, err := redisBroker.open()
	if err != nil {
		return true, fmt.Errorf("Dial: %s", err) // retry true
	}

	defer redisBroker.closeConn(conn)

	psc := redis.PubSubConn{Conn: conn}
	if err := psc.Subscribe(redisBroker.config.DefaultQueue); err != nil {
		return true, err // retry true
	}
	defer redisBroker.closePubSub(psc)

	redisBroker.stopChan = make(chan int)
	deliveries := make(chan signatures.TaskSignature)
	errors := make(chan error)

	log.Print("[*] Waiting for messages. To exit press CTRL+C")

	// Receving goroutine
	go func() {
		for {
			switch n := psc.Receive().(type) {
			case redis.Message:
				log.Printf("Received new message: %s", n.Data)

				signature := signatures.TaskSignature{}
				if err := json.Unmarshal(n.Data, &signature); err != nil {
					errors <- err
					continue
				}

				deliveries <- signature
			case redis.Subscription:
				if n.Count == 0 {
					return
				}
			case error:
				log.Print(n)
			}
		}
	}()

	// Iterate over delivered tasks and process them
	for {
		select {
		case signature := <-deliveries:
			taskProcessor.Process(&signature)
		case err := <-errors:
			return true, err // retry true
		case <-redisBroker.stopChan:
			return false, nil // retry false
		}
	}
}

// StopConsuming quits the loop
func (redisBroker *RedisBroker) StopConsuming() {
	// Notifying the quit channel stops consuming of messages
	redisBroker.stopChan <- 1
}

// Publish places a new message on the default queue
func (redisBroker *RedisBroker) Publish(signature *signatures.TaskSignature) error {
	conn, err := redisBroker.open()
	if err != nil {
		fmt.Errorf("Dial: %s", err)
	}

	defer redisBroker.closeConn(conn)

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON Encode Message: %v", err)
	}

	conn.Do("PUBLISH", redisBroker.config.DefaultQueue, message)
	return conn.Flush()
}

func (redisBroker *RedisBroker) open() (redis.Conn, error) {
	network := "tcp"
	address := strings.Split(redisBroker.config.Broker, "redis://")[1]
	return redis.Dial(network, address)
}

func (redisBroker *RedisBroker) closeConn(conn redis.Conn) error {
	if err := conn.Close(); err != nil {
		return fmt.Errorf("Connection Close: %s", err)
	}

	return nil
}

func (redisBroker *RedisBroker) closePubSub(psc redis.PubSubConn) error {
	// Unsubscribe from all connections. This will cause the receiving
	// goroutine to exit.
	if err := psc.Unsubscribe(); err != nil {
		return fmt.Errorf("PubSub Unsubscribe: %s", err)
	}

	if err := psc.Close(); err != nil {
		return fmt.Errorf("PubSub Close: %s", err)
	}

	return nil
}
{'fraction_non_alphanumeric': 0.12375834554632796, 'fraction_numerical': 0.0016283992835043153, 'mean_word_length': 3.7465224111282844, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 3, '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': '4511880', 'n_tokens_mistral': 2322, 'n_tokens_neox': 2097, 'n_words': 563}