<commit_before>package storage

import (
	"time"

	"encoding/json"
	"fmt"

	"sort"

	"github.com/achilleasa/trace"
	"github.com/garyburd/redigo/redis"
)

// This storage backend is built on top of Redis. Internally it uses
// a connection pool to provide thread-safe access.
type Redis struct {
	connPool *redis.Pool
}

// Create a new Redis storage.
func NewRedis(redisEndpoint string, password string, db uint, timeout time.Duration) *Redis {
	return &Redis{
		connPool: &redis.Pool{
			MaxIdle:     3,
			IdleTimeout: 240 * time.Second,
			Dial: func() (redis.Conn, error) {
				c, err := redis.DialTimeout("tcp", redisEndpoint, timeout, timeout, timeout)
				if err != nil {
					return nil, err
				}
				if password != "" {
					if _, err = c.Do("AUTH", password); err != nil {
						c.Close()
						return nil, err
					}
				}
				if db > 0 {
					if _, err = c.Do("SELECT", db); err != nil {
						c.Close()
						return nil, err
					}
				}

				return c, err
			},
			TestOnBorrow: func(c redis.Conn, t time.Time) error {
				_, err := c.Do("PING")
				return err
			},
		},
	}
}

// Store a trace entry and set a TTL on it. If the ttl is 0 then the
// trace record will never expire. Implements the Storage interface.
func (r *Redis) Store(logEntry *trace.Record, ttl time.Duration) error {
	json, err := json.Marshal(logEntry)
	if err != nil {
		return err
	}

	conn := r.connPool.Get()
	defer conn.Close()

	conn.Send("MULTI")

	// Append log entry to a list that shares the same traceId
	// and set a TTL
	traceKey := fmt.Sprintf("trace.%s", logEntry.TraceId)
	conn.Send("LPUSH", traceKey, json)
	if ttl > time.Second {
		conn.Send("EXPIRE", traceKey, ttl.Seconds())
	}

	// Add logEntry.From to the set of known services
	conn.Send("SADD", "trace.services", logEntry.From)

	// If this is an outgoing request, add the destination to the dependency set
	// for the origin
	if logEntry.Type == trace.Request {
		conn.Send("SADD", fmt.Sprintf("trace.%s.deps", logEntry.From), logEntry.To)
	}

	// Exec pipeline
	_, err = conn.Do("EXEC")
	return err
}

// Fetch a set of time-ordered trace entries with the given trace-id.
func (r *Redis) GetTrace(traceId string) (trace.Trace, error) {

	conn := r.connPool.Get()
	defer conn.Close()

	// Get the number of records
	traceKey := fmt.Sprintf("trace.%s", traceId)
	len, err := redis.Int(conn.Do("LLEN", traceKey))
	if err != nil {
		return nil, err
	}

	// Fetch all records
	rawRows, err := redis.Strings(conn.Do("LRANGE", traceKey, 0, len))
	if err != nil {
		return nil, err
	}

	// Unmarshal raw data
	traceLog := make(trace.Trace, len)
	for index, rawRow := range rawRows {
		entry := trace.Record{}
		err = json.Unmarshal([]byte(rawRow), &entry)
		if err != nil {
			return nil, err
		}
		traceLog[index] = entry
	}

	// Sort the trace so entries appear in insertion order
	sort.Sort(traceLog)

	return traceLog, nil
}

// Get service dependencies optionally filtered by a set of service names. If no filters are
// specified then the response will include all services currently known to the storage.
func (r *Redis) GetDependencies(srvFilter ...string) ([]trace.Dependencies, error) {
	conn := r.connPool.Get()
	defer conn.Close()

	var err error

	if len(srvFilter) == 0 {
		srvFilter, err = redis.Strings(conn.Do("SMEMBERS", "trace.services"))
		if err != nil {
			return nil, err
		}
	}

	// Fetch deps in a single batch
	conn.Send("MULTI")
	for _, serviceName := range srvFilter {
		conn.Send("SMEMBERS", fmt.Sprintf("trace.%s.deps", serviceName))
	}
	replies, err := redis.Values(conn.Do("EXEC"))
	if err != nil {
		return nil, err
	}

	// Assemble deps
	replyCount := len(srvFilter)
	serviceDeps := make([]trace.Dependencies, replyCount)
	for index := 0; index < replyCount; index++ {
		deps, _ := redis.Strings(replies[index], nil)
		serviceDeps[index] = trace.Dependencies{
			Service:      srvFilter[index],
			Dependencies: deps,
		}
	}
<commit_msg>Make sure services are sorted by name<commit_after>package storage

import (
	"time"

	"encoding/json"
	"fmt"

	"sort"

	"github.com/achilleasa/trace"
	"github.com/garyburd/redigo/redis"
)

// This storage backend is built on top of Redis. Internally it uses
// a connection pool to provide thread-safe access.
type Redis struct {
	connPool *redis.Pool
}

// Create a new Redis storage.
func NewRedis(redisEndpoint string, password string, db uint, timeout time.Duration) *Redis {
	return &Redis{
		connPool: &redis.Pool{
			MaxIdle:     3,
			IdleTimeout: 240 * time.Second,
			Dial: func() (redis.Conn, error) {
				c, err := redis.DialTimeout("tcp", redisEndpoint, timeout, timeout, timeout)
				if err != nil {
					return nil, err
				}
				if password != "" {
					if _, err = c.Do("AUTH", password); err != nil {
						c.Close()
						return nil, err
					}
				}
				if db > 0 {
					if _, err = c.Do("SELECT", db); err != nil {
						c.Close()
						return nil, err
					}
				}

				return c, err
			},
			TestOnBorrow: func(c redis.Conn, t time.Time) error {
				_, err := c.Do("PING")
				return err
			},
		},
	}
}

// Store a trace entry and set a TTL on it. If the ttl is 0 then the
// trace record will never expire. Implements the Storage interface.
func (r *Redis) Store(logEntry *trace.Record, ttl time.Duration) error {
	json, err := json.Marshal(logEntry)
	if err != nil {
		return err
	}

	conn := r.connPool.Get()
	defer conn.Close()

	conn.Send("MULTI")

	// Append log entry to a list that shares the same traceId
	// and set a TTL
	traceKey := fmt.Sprintf("trace.%s", logEntry.TraceId)
	conn.Send("LPUSH", traceKey, json)
	if ttl > time.Second {
		conn.Send("EXPIRE", traceKey, ttl.Seconds())
	}

	// Add logEntry.From to the set of known services
	conn.Send("SADD", "trace.services", logEntry.From)

	// If this is an outgoing request, add the destination to the dependency set
	// for the origin
	if logEntry.Type == trace.Request {
		conn.Send("SADD", fmt.Sprintf("trace.%s.deps", logEntry.From), logEntry.To)
	}

	// Exec pipeline
	_, err = conn.Do("EXEC")
	return err
}

// Fetch a set of time-ordered trace entries with the given trace-id.
func (r *Redis) GetTrace(traceId string) (trace.Trace, error) {

	conn := r.connPool.Get()
	defer conn.Close()

	// Get the number of records
	traceKey := fmt.Sprintf("trace.%s", traceId)
	len, err := redis.Int(conn.Do("LLEN", traceKey))
	if err != nil {
		return nil, err
	}

	// Fetch all records
	rawRows, err := redis.Strings(conn.Do("LRANGE", traceKey, 0, len))
	if err != nil {
		return nil, err
	}

	// Unmarshal raw data
	traceLog := make(trace.Trace, len)
	for index, rawRow := range rawRows {
		entry := trace.Record{}
		err = json.Unmarshal([]byte(rawRow), &entry)
		if err != nil {
			return nil, err
		}
		traceLog[index] = entry
	}

	// Sort the trace so entries appear in insertion order
	sort.Sort(traceLog)

	return traceLog, nil
}

// Get service dependencies optionally filtered by a set of service names. If no filters are
// specified then the response will include all services currently known to the storage.
func (r *Redis) GetDependencies(srvFilter ...string) ([]trace.Dependencies, error) {
	conn := r.connPool.Get()
	defer conn.Close()

	var err error

	if len(srvFilter) == 0 {
		srvFilter, err = redis.Strings(conn.Do("SMEMBERS", "trace.services"))
		if err != nil {
			return nil, err
		}
		sort.Strings(srvFilter)
	}

	// Fetch deps in a single batch
	conn.Send("MULTI")
	for _, serviceName := range srvFilter {
		conn.Send("SMEMBERS", fmt.Sprintf("trace.%s.deps", serviceName))
	}
	replies, err := redis.Values(conn.Do("EXEC"))
	if err != nil {
		return nil, err
	}

	// Assemble deps
	replyCount := len(srvFilter)
	serviceDeps := make([]trace.Dependencies, replyCount)
	for index := 0; index < replyCount; index++ {
		deps, _ := redis.Strings(replies[index], nil)
		serviceDeps[index] = trace.Dependencies{
			Service:      srvFilter[index],
			Dependencies: deps,
		}
	}

{'fraction_non_alphanumeric': 0.13701923076923078, 'fraction_numerical': 0.002277327935222672, 'mean_word_length': 3.216, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, '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': '4478714', 'n_tokens_mistral': 3014, 'n_tokens_neox': 2706, 'n_words': 843}