<commit_before>// Package db manages the queuing/persistance for the postmaster
package db

import (
	"encoding/json"
	"errors"
	"time"

	"github.com/levenlabs/go-llog"
	"github.com/levenlabs/golib/genapi"
	"github.com/levenlabs/postmaster/ga"
	"github.com/levenlabs/postmaster/sender"
	"github.com/mediocregopher/okq-go.v2"
)

var (
	normalQueue     = "email-normal"
	statsQueue      = "stats-normal"
	uniqueArgStatID = "pmStatsID"
	uniqueArgEnvID  = "pmEnvID"
)

var jobCh chan job

var useOkq bool

type job struct {
	Queue    string
	Contents string
	RespCh   chan error
}

func init() {
	ga.GA.AppendInit(func(g *genapi.GenAPI) {
		if ga.GA.OkqInfo.Client == nil {
			return
		}
		jobCh = make(chan job)

		okqClient := ga.GA.OkqInfo.Client
		// Receive jobs from StoreSendJob() and StoreStatsJob() and Push into okq
		go func() {
			for job := range jobCh {
				job.RespCh <- okqClient.Push(job.Queue, job.Contents, okq.Normal)
			}
		}()

		// Receive jobs from okq and send to sender
		consumeSpin(handleSendEvent, normalQueue)

		// Receive jobs from okq and store in stats
		consumeSpin(handleStatsEvent, statsQueue)

		useOkq = true
	})
}

// DisableOkq turns off using okq for job storing
// this should ONLY be called during testing
func DisableOkq() {
	useOkq = false
}

func consumeSpin(fn okq.ConsumerFunc, q string) {
	llog.Info("creating okq consumer", llog.KV{"queue": q})
	consumer := ga.GA.OkqInfo.Client
	go func(c *okq.Client) {
		for {
			err := <-c.Consumer(fn, nil, q)
			llog.Error("consumer error", llog.KV{"queue": q}, llog.ErrKV(err))
			time.Sleep(10 * time.Second)
		}
	}(consumer)
}

// StoreSendJob creates a new Mail job with jobContents and sends it to okq
func StoreSendJob(jobContents string) error {
	if !useOkq {
		if !sendEmail(jobContents) {
			return errors.New("Failed to send email (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{normalQueue, jobContents, respCh}
	return <-respCh
}

// StoreStatsJob creates a new statsJob with jobContents and sends it to okq
func StoreStatsJob(jobContents string) error {
	if !useOkq {
		if !storeStats(jobContents) {
			return errors.New("Failed to store stats (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{statsQueue, jobContents, respCh}
	return <-respCh
}

func handleSendEvent(e okq.Event) bool {
	return sendEmail(e.Contents)
}

func handleStatsEvent(e okq.Event) bool {
	return storeStats(e.Contents)
}

func sendEmail(jobContents string) bool {
	job := new(sender.Mail)
	err := json.Unmarshal([]byte(jobContents), job)
	if err != nil {
		llog.Error("error json decoding into sender.Mail", llog.KV{
			"jobContents": jobContents,<commit_msg>fix some backwards incompatible changes to okq driver<commit_after>// Package db manages the queuing/persistance for the postmaster
package db

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	"github.com/levenlabs/go-llog"
	"github.com/levenlabs/golib/genapi"
	"github.com/levenlabs/postmaster/ga"
	"github.com/levenlabs/postmaster/sender"
	"github.com/mediocregopher/okq-go.v2"
)

var (
	normalQueue     = "email-normal"
	statsQueue      = "stats-normal"
	uniqueArgStatID = "pmStatsID"
	uniqueArgEnvID  = "pmEnvID"
)

var jobCh chan job

var useOkq bool

type job struct {
	Queue    string
	Contents string
	RespCh   chan error
}

func init() {
	ga.GA.AppendInit(func(g *genapi.GenAPI) {
		if ga.GA.OkqInfo.Client == nil {
			return
		}
		jobCh = make(chan job)

		okqClient := ga.GA.OkqInfo.Client
		// Receive jobs from StoreSendJob() and StoreStatsJob() and Push into okq
		go func() {
			for job := range jobCh {
				job.RespCh <- okqClient.Push(job.Queue, job.Contents, okq.Normal)
			}
		}()

		// Receive jobs from okq and send to sender
		consumeSpin(handleSendEvent, normalQueue)

		// Receive jobs from okq and store in stats
		consumeSpin(handleStatsEvent, statsQueue)

		useOkq = true
	})
}

// DisableOkq turns off using okq for job storing
// this should ONLY be called during testing
func DisableOkq() {
	useOkq = false
}

func consumeSpin(fn okq.ConsumerFunc, q string) {
	llog.Info("creating okq consumer", llog.KV{"queue": q})
	consumer := ga.GA.OkqInfo.Client
	go func(c *okq.Client) {
		for {
			err := <-c.Consumer(fn, nil, q)
			llog.Error("consumer error", llog.KV{"queue": q}, llog.ErrKV(err))
			time.Sleep(10 * time.Second)
		}
	}(consumer)
}

// StoreSendJob creates a new Mail job with jobContents and sends it to okq
func StoreSendJob(jobContents string) error {
	if !useOkq {
		if !sendEmail(jobContents) {
			return errors.New("Failed to send email (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{normalQueue, jobContents, respCh}
	return <-respCh
}

// StoreStatsJob creates a new statsJob with jobContents and sends it to okq
func StoreStatsJob(jobContents string) error {
	if !useOkq {
		if !storeStats(jobContents) {
			return errors.New("Failed to store stats (bypassing okq)")
		}
		return nil
	}
	respCh := make(chan error)
	jobCh <- job{statsQueue, jobContents, respCh}
	return <-respCh
}

func handleSendEvent(_ context.Context, e okq.Event) bool {
	return sendEmail(e.Contents)
}

func handleStatsEvent(_ context.Context, e okq.Event) bool {
	return storeStats(e.Contents)
}

func sendEmail(jobContents string) bool {
	job := new(sender.Mail)
	err := json.Unmarshal([]byte(jobContents), job)
	if err != nil {
		llog.Error("error json decoding into sender.Mail", llog.KV{
			"jobContents": jobContents,
{'fraction_non_alphanumeric': 0.11717466691002007, 'fraction_numerical': 0.0010950903449534588, 'mean_word_length': 3.8841354723707666, 'pattern_counts': {'":': 6, '<': 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': '5228720', 'n_tokens_mistral': 2099, 'n_tokens_neox': 1953, 'n_words': 525}