llmongodb/consumer.go

// consumer
package main

import (
	"fmt"
	//	"github.com/garyburd/redigo/redis"
	"crypto/sha256"
	"encoding/hex"
	"log"
	"strconv"
	"strings"
	"time"

	"gopkg.in/mgo.v2"
)

type consumed struct {
	user   string
	error  bool
	logins []string
	empty  bool
}

func hash(val []byte) string {

	h := sha256.New()
	h.Write(val)

	return hex.EncodeToString(h.Sum(nil))
}

func consumer() {

	for {

		prod := <-consume

		status = _Consumer

		var bulk = make(map[string]*mgo.Bulk)
		var rtbulk []MongoLogin
		var col = make(map[string]*mgo.Collection)
		var slogin = make(map[string][]string)

		if opts.Bulk {
			bulk[opts.Month] = dbs.ll.Bulk()
			bulk[opts.Month].Unordered()
		} else {
			col[opts.Month] = dbs.ll
		}

		cons := consumed{
			user:   prod.user,
			logins: make([]string, 0),
			error:  false,
			empty:  true,
		}

		start := time.Now()
		for i := range prod.logins {
			login := prod.logins[i]
			// se la riga di login e' vuota
			if login == "" {
				log.Println("Login empty: ", prod.user)
				cons.logins = append(cons.logins, login)
				continue
			}
			sval := strings.Split(login, ":")
			// se il formato della riga di login non e' corretto
			if sval[1] == "" {
				log.Println("Login format error: ", login, prod.user)
				cons.logins = append(cons.logins, login)
				continue
			}
			// se il timestamp della riga di login non e' corretto
			date, err := strconv.ParseInt(sval[1], 10, 64)
			if err != nil {
				log.Printf("Date Error: %+v - %s - %s\n", err, prod.user, login)
				cons.logins = append(cons.logins, login)
				continue
			}
			ml := MongoLogin{
				// genera l' _ID con user e timestamp
				ID:       hash([]byte(fmt.Sprintf("%s%s", prod.user, time.Unix(date, 0).Format("20060102T1504")))), // Format("20060102T150405")
				User:     prod.user,
				Protocol: sval[0],
				IP:       sval[2],
				Date:     time.Unix(date, 0),
				Insert:   time.Now(),
			}

			if opts.Month != ml.Date.Format("0601") {
				dt := fmt.Sprintf("lastlogin_%s", ml.Date.Format("0601"))
				if opts.Bulk {
					if dbs.isMongodb() {
						if _, ok := bulk[dt]; !ok {
							bulk[dt] = dbs.mdb.DB("lastlogin").C(dt).Bulk()
							bulk[dt].Unordered()
						}
						bulk[dt].Insert(ml)
						slogin[dt] = append(slogin[dt], login)
					}
					if dbs.isRethink() {
						rtbulk = append(rtbulk, ml)
						slogin["rt"] = append(slogin["rt"], login)
					}
				} else {
					if dbs.isMongodb() {
						if _, ok := col[dt]; !ok {
							col[dt] = dbs.mdb.DB("lastlogin").C(dt)
						}
						err = col[dt].Insert(ml)
						if err != nil {
							if !strings.Contains(err.Error(), "E11000") {
								fmt.Printf("Mongo Insert Err: %+v\n", err)
								cons.error = true
								counter <- Counterchan{
									tipo: "err",
									val:  1,
								}
								continue
							} else {
								counter <- Counterchan{
									tipo: "dup",
									val:  1,
								}
							}
						}
					}
					if dbs.isRethink() {
						_, err = dbs.rtdb.Insert(ml)
						if err != nil {
							if !strings.Contains(err.Error(), "Duplicate primary key") {
								fmt.Printf("RT Insert Err: %+v\n", err)
								cons.error = true
								counter <- Counterchan{
									tipo: "err",
									val:  1,
								}
								continue
							} else {
								counter <- Counterchan{
									tipo: "dup",
									val:  1,
								}
							}
						}
					}

					cons.logins = append(cons.logins, login)
				}
			} else {
				// inserisce il login su Mongodb
				if opts.Bulk {
					if dbs.isMongodb() {
						bulk[opts.Month].Insert(ml)
						slogin[opts.Month] = append(slogin[opts.Month], login)
					}
					if dbs.isRethink() {
						rtbulk = append(rtbulk, ml)
						slogin["rt"] = append(slogin["rt"], login)
					}
				} else {
					if dbs.isMongodb() {
						err = col[opts.Month].Insert(ml)
						if err != nil {
							if !strings.Contains(err.Error(), "E11000") {
								fmt.Printf("Err: %+v\n", err)
								cons.error = true
								counter <- Counterchan{
									tipo: "err",
									val:  1,
								}
								continue
							} else {
								counter <- Counterchan{
									tipo: "dup",
									val:  1,
								}
							}
						}
					}
					if dbs.isRethink() {
						resp, err := dbs.rtdb.Insert(ml)
						counter <- Counterchan{
							tipo: "ins",
							val:  resp.Inserted,
						}
						if err != nil {
							if !strings.Contains(err.Error(), "Duplicate primary key") {
								fmt.Printf("RT Insert Err: %+v\n", err)
								cons.error = true
								counter <- Counterchan{
									tipo: "err",
									val:  1,
								}
								continue
							} else {
								counter <- Counterchan{
									tipo: "dup",
									val:  1,
								}
							}
						}
					}

					cons.logins = append(cons.logins, login)
				}
			}
		}

		if opts.Bulk {
			if dbs.isMongodb() {
				for key, _ := range bulk {
					_, err := bulk[key].Run()
					if err != nil {
						if !strings.Contains(err.Error(), "E11000") {
							fmt.Printf("Err: %+v\n", err)
							cons.error = true
							counter <- Counterchan{
								tipo: "err",
								val:  len(slogin[key]),
							}
							continue
						} else {
							counter <- Counterchan{
								tipo: "dup",
								val:  strings.Count(err.Error(), "E11000"),
							}
						}
					}
					cons.logins = append(cons.logins, slogin[key]...)
				}
			}
			if dbs.isRethink() {
				resp, err := dbs.rtdb.MultiInsert(rtbulk)
				counter <- Counterchan{
					tipo: "ins",
					val:  resp.Inserted,
				}
				if err != nil {
					if !strings.Contains(err.Error(), "Duplicate primary key") {
						cons.error = true
						counter <- Counterchan{
							tipo: "err",
							val:  1,
						}
						continue
					} else {
						counter <- Counterchan{
							tipo: "dup",
							val:  1,
						}
					}
				}
				cons.logins = append(cons.logins, slogin["rt"]...)
			}
		}

		counter <- Counterchan{
			tipo: "log",
			val:  len(prod.logins),
		}

		if opts.MaxLogins > -1 && len(prod.logins) < opts.MaxLogins {
			cons.empty = false
		}

		if opts.Debug {
			fmt.Printf("CONS: user=%s logins=%d in %v - active=%d\n", prod.user, len(prod.logins), time.Since(start), dbs.rdb.ActiveCount())
		}

		//		wg.Done()
		remove <- cons
	}
}