<commit_before>package main

import (
	"encoding/binary"
	nsq "github.com/bitly/go-nsq"
	"github.com/boltdb/bolt"
	log "github.com/gonet2/libs/nsq-logger"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

const (
	DEFAULT_NSQLOOKUPD   = "http://172.17.42.1:4161"
	ENV_NSQLOOKUPD       = "NSQLOOKUPD_HOST"
	TOPIC                = "REDOLOG"
	CHANNEL              = "ARCH"
	SERVICE              = "[ARCH]"
	REDO_TIME_FORMAT     = "REDO-2006-01-02T15:04:05.RDO"
	REDO_ROTATE_INTERVAL = 24 * time.Hour
	BOLTDB_BUCKET        = "REDOLOG"
	DATA_DIRECTORY       = "/data/"
)

type Archiver struct {
	pending chan []byte
	stop    chan bool
}

func (arch *Archiver) init() {
	arch.pending = make(chan []byte)
	arch.stop = make(chan bool)
	cfg := nsq.NewConfig()
	consumer, err := nsq.NewConsumer(TOPIC, CHANNEL, cfg)
	if err != nil {
		log.Critical(err)
		os.Exit(-1)
	}

	// message process
	consumer.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
		arch.pending <- msg.Body
		return nil
	}))

	// read environtment variable
	addresses := []string{DEFAULT_NSQLOOKUPD}
	if env := os.Getenv(ENV_NSQLOOKUPD); env != "" {
		addresses = strings.Split(env, ";")
	}

	// connect to nsqlookupd
	log.Trace("connect to nsqlookupds ip:", addresses)
	if err := consumer.ConnectToNSQLookupds(addresses); err != nil {
		log.Critical(err)
		return
	}
	log.Info("nsqlookupd connected")

	go arch.archive_task()
}

func (arch *Archiver) archive_task() {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGTERM)
	timer := time.After(REDO_ROTATE_INTERVAL)
	db := arch.new_redolog()
	key := make([]byte, 8)
	for {
		select {
		case msg := <-arch.pending:
			db.Update(func(tx *bolt.Tx) error {
				b := tx.Bucket([]byte(BOLTDB_BUCKET))
				id, err := b.NextSequence()
				if err != nil {
					log.Critical(err)
					return err
				}
				binary.BigEndian.PutUint64(key, uint64(id))
				if err = b.Put(key, msg); err != nil {
					log.Critical(err)
					return err
				}
				return nil
			})
		case <-timer:
			db.Close()
			// rotate redolog
			db = arch.new_redolog()
			timer = time.After(REDO_ROTATE_INTERVAL)
		case <-sig:
			db.Close()
			log.Info("SIGTERM")
			os.Exit(0)
		}
	}
}

func (arch *Archiver) new_redolog() *bolt.DB {
	file := DATA_DIRECTORY + time.Now().Format(REDO_TIME_FORMAT)
	log.Info(file)
	db, err := bolt.Open(file, 0600, nil)
	if err != nil {
		log.Critical(err)
		os.Exit(-1)
	}
	// create bulket<commit_msg>batch insert into boltdb<commit_after>package main

import (
	"encoding/binary"
	nsq "github.com/bitly/go-nsq"
	"github.com/boltdb/bolt"
	log "github.com/gonet2/libs/nsq-logger"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

const (
	DEFAULT_NSQLOOKUPD   = "http://172.17.42.1:4161"
	ENV_NSQLOOKUPD       = "NSQLOOKUPD_HOST"
	TOPIC                = "REDOLOG"
	CHANNEL              = "ARCH"
	SERVICE              = "[ARCH]"
	REDO_TIME_FORMAT     = "REDO-2006-01-02T15:04:05.RDO"
	REDO_ROTATE_INTERVAL = 24 * time.Hour
	BOLTDB_BUCKET        = "REDOLOG"
	DATA_DIRECTORY       = "/data/"
	BATCH_SIZE           = 1024
	SYNC_INTERVAL        = 10 * time.Millisecond
)

type Archiver struct {
	pending chan []byte
	stop    chan bool
}

func (arch *Archiver) init() {
	arch.pending = make(chan []byte, BATCH_SIZE)
	arch.stop = make(chan bool)
	cfg := nsq.NewConfig()
	consumer, err := nsq.NewConsumer(TOPIC, CHANNEL, cfg)
	if err != nil {
		log.Critical(err)
		os.Exit(-1)
	}

	// message process
	consumer.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
		arch.pending <- msg.Body
		return nil
	}))

	// read environtment variable
	addresses := []string{DEFAULT_NSQLOOKUPD}
	if env := os.Getenv(ENV_NSQLOOKUPD); env != "" {
		addresses = strings.Split(env, ";")
	}

	// connect to nsqlookupd
	log.Trace("connect to nsqlookupds ip:", addresses)
	if err := consumer.ConnectToNSQLookupds(addresses); err != nil {
		log.Critical(err)
		return
	}
	log.Info("nsqlookupd connected")

	go arch.archive_task()
}

func (arch *Archiver) archive_task() {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGTERM)
	timer := time.After(REDO_ROTATE_INTERVAL)
	sync_ticker := time.NewTicker(SYNC_INTERVAL)
	db := arch.new_redolog()
	key := make([]byte, 8)
	for {
		select {
		case <-sync_ticker.C:
			n := len(arch.pending)
			if n == 0 {
				continue
			}

			db.Update(func(tx *bolt.Tx) error {
				b := tx.Bucket([]byte(BOLTDB_BUCKET))
				for i := 0; i < n; i++ {
					id, err := b.NextSequence()
					if err != nil {
						log.Critical(err)
						continue
					}
					binary.BigEndian.PutUint64(key, uint64(id))
					if err = b.Put(key, <-arch.pending); err != nil {
						log.Critical(err)
						continue
					}
				}
				return nil
			})
		case <-timer:
			db.Close()
			// rotate redolog
			db = arch.new_redolog()
			timer = time.After(REDO_ROTATE_INTERVAL)
		case <-sig:
			db.Close()
			log.Info("SIGTERM")
			os.Exit(0)
		}
	}
}

func (arch *Archiver) new_redolog() *bolt.DB {
	file := DATA_DIRECTORY + time.Now().Format(REDO_TIME_FORMAT)
	log.Info(file)
	db, err := bolt.Open(file, 0600, nil)
	if err != nil {
		log.Critical(err)
		os.Exit(-1)
	}
	// create bulket
{'fraction_non_alphanumeric': 0.1509397024275646, 'fraction_numerical': 0.018010963194988253, 'mean_word_length': 3.0323599052880823, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, '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': '4637550', 'n_tokens_mistral': 2325, 'n_tokens_neox': 2025, 'n_words': 395}