<commit_before>/*
    The MIT License (MIT)

	Copyright (c) 2015 myhug.cn and zhouwench (zhouwench@gmail.com)

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.
*/
package ktransfer

import (
		"github.com/dzch/go-utils/logger"
		"github.com/dzch/kafka/consumergroup"
		"github.com/tinylib/msgp/msgp"
		"gopkg.in/Shopify/sarama.v1"
		"hash/fnv"
		"time"
		"fmt"
		"errors"
		"hash"
		"bytes"
		"math/rand"
	   )

const (
		WAIT_ACK_NONBLOCK = iota
		WAIT_ACK_BLOCK
		CHECK_NEED_WAIT_ACK
		WAIT_ACK_AND_DI
	  )

var (
		gDefaultMethod = "*"
	)

type TransData struct {
	topic string
	partition int32
	method string
	key string
	transid int64
	data []byte
}

type AckData struct {
	workerId uint32
	transid int64
}

type TransDi struct {
	moduleName string
	zkHosts []string
	zkChroot string
	zkTimeout time.Duration
	ackedMinTransid int64
	curReadTransid int64 // 已经读取的最大transid
	curDiTransid int64
	protocolConfig map[interface{}]interface{}
	backendServers []string
	windowSize int
	workerNum uint32
	transWorkers []*TransWorker
	fatalErrorChan chan error
	state int
	ackDataChan chan *AckData
	transWindow *TransWindow
	topic string
	methodEnabled map[string]bool
	maxRetryTimes int    // -1, 无限重试
	failRetryInterval time.Duration
	cg *consumergroup.ConsumerGroup
	dispatchHasher hash.Hash32
}

func (td *TransDi) init() (err error) {
	err = td.initChans()
	if err != nil {
		return
	}
	err = td.initTransWindow()
	if err != nil {
		return
	}
	err = td.initTransWorkers()
	if err != nil {
		return
	}
	err = td.initConsumer()
	if err != nil {
		return
	}
	err = td.initDispatchHash()
	if err != nil {
		return
	}
	return nil
}

func (td *TransDi) run() {
	logger.Debug("transdi start for module [%s]", td.moduleName)
	for _, worker := range td.transWorkers {
		go worker.run()
	}
	td.state = WAIT_ACK_NONBLOCK
	for {
		switch td.state {
			case WAIT_ACK_NONBLOCK:
				td.waitAckNonBlock()
				continue
			case CHECK_NEED_WAIT_ACK:
				td.checkNeedWaitAck()
				continue
			case WAIT_ACK_BLOCK:
				td.waitAckBlock()
				continue
			case WAIT_ACK_AND_DI:
				td.waitAckAndDi()
				continue
			default:
				td.fatalErrorChan <-errors.New(fmt.Sprintf("invalid transdi run state: %d", td.state))
				break
		}
	}
}

func (td *TransDi) initChans() (err error) {
	td.ackDataChan = make(chan *AckData)
	return nil
}

func (td *TransDi) initTransWindow() (err error) {
	td.transWindow = &TransWindow {
        windowSize: td.windowSize,
	}
	return td.transWindow.init()
}

func (td *TransDi) initConsumer() (err error) {
    config := consumergroup.NewConfig()
	config.Zookeeper.Chroot = td.zkChroot
	config.Zookeeper.Timeout = td.zkTimeout
	td.cg, err = consumergroup.JoinConsumerGroup(
			td.moduleName,
			[]string{td.topic},
			td.zkHosts,
			config)
	if err != nil {
		return err
	}
	return nil
}

func (td *TransDi) initDispatchHash() (err error) {
	td.dispatchHasher = fnv.New32a()
	return nil
}

func (td *TransDi) initTransWorkers() (err error) {
    for id := uint32(0); id < td.workerNum; id ++ {
        worker := &TransWorker {
			moduleName: td.moduleName,
            workerId: id,
			backendServers: td.backendServers,
			protocolConfig: td.protocolConfig,
			fatalErrorChan: td.fatalErrorChan,
			inWork: false,
			ackDataChan: td.ackDataChan,
			maxRetryTimes: td.maxRetryTimes,
			failRetryInterval: td.failRetryInterval,
		}
		err = worker.init()
		if err != nil {
			return
		}
		td.transWorkers = append(td.transWorkers, worker)
	}
	logger.Notice("all workers have been inited for moudle [%s]", td.moduleName)
	return nil
}

func (td *TransDi) waitAckNonBlock () {
//	logger.Debug("in waitAckNonBlock")
	for {
		select {
			case ackData := <-td.ackDataChan:
				td.processAckData(ackData)
				/* do not break, so we can receive ack untile no ack, then state = CHECK_NEED_WAIT_ACK */
			default:
				td.state = CHECK_NEED_WAIT_ACK
				return
		}
	}
}

func (td *TransDi) waitAckBlock() {
//	logger.Debug("in waitAckBlock")
	for ackData := range td.ackDataChan {
		td.processAckData(ackData)
		td.state = WAIT_ACK_NONBLOCK
		break
	}
}

func (td *TransDi) checkNeedWaitAck() {
	if td.transWindow.isFull() {
		td.state = WAIT_ACK_BLOCK
	} else {
		td.state = WAIT_ACK_AND_DI
	}
}

func (td *TransDi) waitAckAndDi() {
//	logger.Debug("in waitAckAndDi")
	var ackData *AckData
	for {
		select {
			case ackData =<-td.ackDataChan:
                td.processAckData(ackData)
	            td.state = WAIT_ACK_NONBLOCK
				return
			case message := <-td.cg.Messages():
				td.processConsumerMessage(message)
				return
			case err := <-td.cg.Errors():
				td.processConsumerError(err)
				return
		}
	}
	return
}

func (td *TransDi) processConsumerMessage(msg *sarama.ConsumerMessage) {
	/* unpack */
	/* pack {
	      'topic' => string,
		  'method' => string,
		  'key' => string,
		  'data' => []byte,
		  }
	*/
	// FIXME: need be optimized here
    sleep := 100*time.Millisecond
	for {
	    var err error
		method := gDefaultMethod
		if len(td.methodEnabled) > 0 {
			// if methods is configured, we should check if this method needed be send
	        buf := bytes.NewReader(msg.Value)
	        buf.Seek(0, 0)
	        msgr := msgp.NewReader(buf)
	        dii, err := msgr.ReadIntf()
	        if err != nil {
			    err = errors.New(fmt.Sprintf("fail to de-msgpack: %s", err.Error()))
	            logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
		        time.Sleep(sleep)
		        continue
	        }
	        di, ok := dii.(map[string]interface{})
	        if !ok {
		        err = errors.New("invalid di: should be map")
	            logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
		        time.Sleep(sleep)
		        continue
	        }
	        methodi, ok := di["method"]
	        if !ok {
	            err = errors.New(fmt.Sprintf("invalid di: method not exists"))
	            logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
		        time.Sleep(sleep)
		        continue
	        }
		    method, ok = methodi.(string)
		    if !ok {
			    err = errors.New(fmt.Sprintf("invlid di: method is not string"))
	            logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
		        time.Sleep(sleep)
		        continue
		    }
	        if _, ok = td.methodEnabled[method]; !ok {
	            // not ours
	            return
	        }
		}
	    // dispatch
	    td.transWindow.addSlot(msg)
		var workerId uint32
		if msg.Key == nil || len(msg.Key) == 0 {
			workerId = uint32(rand.Int31n(int32(td.workerNum)))
		} else {
	        td.dispatchHasher.Reset()
	        _, err = td.dispatchHasher.Write(msg.Key)
	        if err != nil {
		        err = errors.New(fmt.Sprintf("fail to compute key-hash: %s", err.Error()))
	            logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
		        time.Sleep(sleep)
		        continue
	        }
            hash := td.dispatchHasher.Sum32()
            workerId = hash % td.workerNum
		}
	    worker := td.transWorkers[workerId]
	    transData := &TransData {
            transid: msg.Offset,
		    topic: msg.Topic,
		    method: method,
			partition: msg.Partition,
		    data: msg.Value,
	    }
	    err = worker.addTrans(transData)
	    if err != nil {
	        logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
		    time.Sleep(sleep)
		    continue
	    }
		td.state = WAIT_ACK_NONBLOCK
		return
	}
}

func (td *TransDi) processConsumerError(cerr *sarama.ConsumerError) {
	logger.Fatal("consumer error: topic=%s, partition=%d, %s", cerr.Topic, cerr.Partition, cerr.Error())
	td.fatalErrorChan <-cerr
}

/* do not change state in this func */
func (td *TransDi) processAckData(ackData *AckData) {
    transid := ackData.transid
	workerId := ackData.workerId
	logger.Debug("process ack: transid=%d, workerId=%d", transid, workerId)
	/* narrow window */
	msg := td.transWindow.ackOne(transid)
	if msg != nil && msg.Offset > td.ackedMinTransid {
		td.ackedMinTransid = msg.Offset
		err := td.cg.CommitUpto(msg)
		if err != nil {
			// TODO: optimized
			logger.Warning("fail to consumergroup.CommitUpto(): %s", err.Error())
			td.fatalErrorChan <-err
			return
		}
		logger.Debug("consumergroup.CommitUpTo %d", msg.Offset)
	}
	logger.Debug("transWindow size: %d", td.transWindow.window.Len())
	/* move worker */
	worker := td.transWorkers[workerId]
	worker.inWork = false
	err := worker.workIfNeed()
	if err != nil {
		logger.Fatal("fail to let worker to work: %s", err.Error())
		td.fatalErrorChan <- err
		return
	}
}

<commit_msg>change Iitial to NewestOffset<commit_after>/*
    The MIT License (MIT)

	Copyright (c) 2015 myhug.cn and zhouwench (zhouwench@gmail.com)

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.
*/
package ktransfer

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/dzch/go-utils/logger"
	"github.com/dzch/kafka/consumergroup"
	"github.com/tinylib/msgp/msgp"
	"gopkg.in/Shopify/sarama.v1"
	"hash"
	"hash/fnv"
	"math/rand"
	"time"
)

const (
	WAIT_ACK_NONBLOCK = iota
	WAIT_ACK_BLOCK
	CHECK_NEED_WAIT_ACK
	WAIT_ACK_AND_DI
)

var (
	gDefaultMethod = "*"
)

type TransData struct {
	topic     string
	partition int32
	method    string
	key       string
	transid   int64
	data      []byte
}

type AckData struct {
	workerId uint32
	transid  int64
}

type TransDi struct {
	moduleName        string
	zkHosts           []string
	zkChroot          string
	zkTimeout         time.Duration
	ackedMinTransid   int64
	curReadTransid    int64 // 已经读取的最大transid
	curDiTransid      int64
	protocolConfig    map[interface{}]interface{}
	backendServers    []string
	windowSize        int
	workerNum         uint32
	transWorkers      []*TransWorker
	fatalErrorChan    chan error
	state             int
	ackDataChan       chan *AckData
	transWindow       *TransWindow
	topic             string
	methodEnabled     map[string]bool
	maxRetryTimes     int // -1, 无限重试
	failRetryInterval time.Duration
	cg                *consumergroup.ConsumerGroup
	dispatchHasher    hash.Hash32
}

func (td *TransDi) init() (err error) {
	err = td.initChans()
	if err != nil {
		return
	}
	err = td.initTransWindow()
	if err != nil {
		return
	}
	err = td.initTransWorkers()
	if err != nil {
		return
	}
	err = td.initConsumer()
	if err != nil {
		return
	}
	err = td.initDispatchHash()
	if err != nil {
		return
	}
	return nil
}

func (td *TransDi) run() {
	logger.Debug("transdi start for module [%s]", td.moduleName)
	for _, worker := range td.transWorkers {
		go worker.run()
	}
	td.state = WAIT_ACK_NONBLOCK
	for {
		switch td.state {
		case WAIT_ACK_NONBLOCK:
			td.waitAckNonBlock()
			continue
		case CHECK_NEED_WAIT_ACK:
			td.checkNeedWaitAck()
			continue
		case WAIT_ACK_BLOCK:
			td.waitAckBlock()
			continue
		case WAIT_ACK_AND_DI:
			td.waitAckAndDi()
			continue
		default:
			td.fatalErrorChan <- errors.New(fmt.Sprintf("invalid transdi run state: %d", td.state))
			break
		}
	}
}

func (td *TransDi) initChans() (err error) {
	td.ackDataChan = make(chan *AckData)
	return nil
}

func (td *TransDi) initTransWindow() (err error) {
	td.transWindow = &TransWindow{
		windowSize: td.windowSize,
	}
	return td.transWindow.init()
}

func (td *TransDi) initConsumer() (err error) {
	config := consumergroup.NewConfig()
	config.Zookeeper.Chroot = td.zkChroot
	config.Zookeeper.Timeout = td.zkTimeout
	config.Offsets.Initial = sarama.OffsetNewest // 如果没有记录，则从最新位置开始消费
	td.cg, err = consumergroup.JoinConsumerGroup(
		td.moduleName,
		[]string{td.topic},
		td.zkHosts,
		config)
	if err != nil {
		return err
	}
	return nil
}

func (td *TransDi) initDispatchHash() (err error) {
	td.dispatchHasher = fnv.New32a()
	return nil
}

func (td *TransDi) initTransWorkers() (err error) {
	for id := uint32(0); id < td.workerNum; id++ {
		worker := &TransWorker{
			moduleName:        td.moduleName,
			workerId:          id,
			backendServers:    td.backendServers,
			protocolConfig:    td.protocolConfig,
			fatalErrorChan:    td.fatalErrorChan,
			inWork:            false,
			ackDataChan:       td.ackDataChan,
			maxRetryTimes:     td.maxRetryTimes,
			failRetryInterval: td.failRetryInterval,
		}
		err = worker.init()
		if err != nil {
			return
		}
		td.transWorkers = append(td.transWorkers, worker)
	}
	logger.Notice("all workers have been inited for moudle [%s]", td.moduleName)
	return nil
}

func (td *TransDi) waitAckNonBlock() {
	//	logger.Debug("in waitAckNonBlock")
	for {
		select {
		case ackData := <-td.ackDataChan:
			td.processAckData(ackData)
			/* do not break, so we can receive ack untile no ack, then state = CHECK_NEED_WAIT_ACK */
		default:
			td.state = CHECK_NEED_WAIT_ACK
			return
		}
	}
}

func (td *TransDi) waitAckBlock() {
	//	logger.Debug("in waitAckBlock")
	for ackData := range td.ackDataChan {
		td.processAckData(ackData)
		td.state = WAIT_ACK_NONBLOCK
		break
	}
}

func (td *TransDi) checkNeedWaitAck() {
	if td.transWindow.isFull() {
		td.state = WAIT_ACK_BLOCK
	} else {
		td.state = WAIT_ACK_AND_DI
	}
}

func (td *TransDi) waitAckAndDi() {
	//	logger.Debug("in waitAckAndDi")
	var ackData *AckData
	for {
		select {
		case ackData = <-td.ackDataChan:
			td.processAckData(ackData)
			td.state = WAIT_ACK_NONBLOCK
			return
		case message := <-td.cg.Messages():
			td.processConsumerMessage(message)
			return
		case err := <-td.cg.Errors():
			td.processConsumerError(err)
			return
		}
	}
	return
}

func (td *TransDi) processConsumerMessage(msg *sarama.ConsumerMessage) {
	/* unpack */
	/* pack {
	      'topic' => string,
		  'method' => string,
		  'key' => string,
		  'data' => []byte,
		  }
	*/
	// FIXME: need be optimized here
	sleep := 100 * time.Millisecond
	for {
		var err error
		method := gDefaultMethod
		if len(td.methodEnabled) > 0 {
			// if methods is configured, we should check if this method needed be send
			buf := bytes.NewReader(msg.Value)
			buf.Seek(0, 0)
			msgr := msgp.NewReader(buf)
			dii, err := msgr.ReadIntf()
			if err != nil {
				err = errors.New(fmt.Sprintf("fail to de-msgpack: %s", err.Error()))
				logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
				time.Sleep(sleep)
				continue
			}
			di, ok := dii.(map[string]interface{})
			if !ok {
				err = errors.New("invalid di: should be map")
				logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
				time.Sleep(sleep)
				continue
			}
			methodi, ok := di["method"]
			if !ok {
				err = errors.New(fmt.Sprintf("invalid di: method not exists"))
				logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
				time.Sleep(sleep)
				continue
			}
			method, ok = methodi.(string)
			if !ok {
				err = errors.New(fmt.Sprintf("invlid di: method is not string"))
				logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
				time.Sleep(sleep)
				continue
			}
			if _, ok = td.methodEnabled[method]; !ok {
				// not ours
				return
			}
		}
		// dispatch
		td.transWindow.addSlot(msg)
		var workerId uint32
		if msg.Key == nil || len(msg.Key) == 0 {
			workerId = uint32(rand.Int31n(int32(td.workerNum)))
		} else {
			td.dispatchHasher.Reset()
			_, err = td.dispatchHasher.Write(msg.Key)
			if err != nil {
				err = errors.New(fmt.Sprintf("fail to compute key-hash: %s", err.Error()))
				logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
				time.Sleep(sleep)
				continue
			}
			hash := td.dispatchHasher.Sum32()
			workerId = hash % td.workerNum
		}
		worker := td.transWorkers[workerId]
		transData := &TransData{
			transid:   msg.Offset,
			topic:     msg.Topic,
			method:    method,
			partition: msg.Partition,
			data:      msg.Value,
		}
		err = worker.addTrans(transData)
		if err != nil {
			logger.Warning("fail to processConsumerMessage: topic=%s, partition=%d, transid=%d, error: %s", msg.Topic, msg.Partition, msg.Offset, err.Error())
			time.Sleep(sleep)
			continue
		}
		td.state = WAIT_ACK_NONBLOCK
		return
	}
}

func (td *TransDi) processConsumerError(cerr *sarama.ConsumerError) {
	logger.Fatal("consumer error: topic=%s, partition=%d, %s", cerr.Topic, cerr.Partition, cerr.Error())
	td.fatalErrorChan <- cerr
}

/* do not change state in this func */
func (td *TransDi) processAckData(ackData *AckData) {
	transid := ackData.transid
	workerId := ackData.workerId
	logger.Debug("process ack: transid=%d, workerId=%d", transid, workerId)
	/* narrow window */
	msg := td.transWindow.ackOne(transid)
	if msg != nil && msg.Offset > td.ackedMinTransid {
		td.ackedMinTransid = msg.Offset
		err := td.cg.CommitUpto(msg)
		if err != nil {
			// TODO: optimized
			logger.Warning("fail to consumergroup.CommitUpto(): %s", err.Error())
			td.fatalErrorChan <- err
			return
		}
		logger.Debug("consumergroup.CommitUpTo %d", msg.Offset)
	}
	logger.Debug("transWindow size: %d", td.transWindow.window.Len())
	/* move worker */
	worker := td.transWorkers[workerId]
	worker.inWork = false
	err := worker.workIfNeed()
	if err != nil {
		logger.Fatal("fail to let worker to work: %s", err.Error())
		td.fatalErrorChan <- err
		return
	}
}
<|endoftext|>


<commit_before>package net

import (
	"io"
	"log"
	"net"
	"sync"

	"../redlot"
)

var info struct {
	sync.RWMutex
	ConnCounter uint64
	TotalCalls  uint64
	reply       Reply
}

func init() {
	// Register commands.
	// system info
	REG("INFO", STATUS_REPLY, Info)

	// KV type
	REG("GET", BULK_REPLY, redlot.Get)
	REG("SET", STATUS_REPLY, redlot.Set)
	REG("DEL", STATUS_REPLY, redlot.Del)
	REG("EXISTS", INT_REPLY, redlot.Exists)
	REG("SETX", STATUS_REPLY, redlot.Setx)
	REG("SETEX", STATUS_REPLY, redlot.Setx) // Alias of SETX
	REG("TTL", INT_REPLY, redlot.Ttl)
	REG("EXPIRE", INT_REPLY, redlot.Expire)

}

func Serve(addr string, options *redlot.Options) {
	// Open LevelDB with options.
	redlot.Open(options)

	// Create sockets listener.
	l, err := net.Listen("tcp4", addr)
	if err != nil {
		log.Fatalf("Listen error: %v\n", err.Error())
	}
	defer l.Close()

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Fatalf("Wait for a connection error: %s\n", err.Error())
		}

		// Count connecion
		info.Lock()
		info.ConnCounter++
		info.Unlock()

		go func(c net.Conn) {
			for {
				req, err := newRequset(c)
				if err == io.EOF {
					break
				}
				if err != nil {
					continue
				}

				info.Lock()
				info.TotalCalls++
				info.Unlock()

				reply := RUN(req.Cmd, req.Args)
				reply.WriteTo(c)
			}

			c.Close()
			info.Lock()
			info.ConnCounter--
			info.Unlock()

		}(conn)

	}
}
<commit_msg>Register  KEYS and SCAN commands.<commit_after>package net

import (
	"io"
	"log"
	"net"
	"sync"

	"../redlot"
)

var info struct {
	sync.RWMutex
	ConnCounter uint64
	TotalCalls  uint64
	reply       Reply
}

func init() {
	// Register commands.
	// system info
	REG("INFO", STATUS_REPLY, Info)

	// KV type
	REG("GET", BULK_REPLY, redlot.Get)
	REG("SET", STATUS_REPLY, redlot.Set)
	REG("DEL", STATUS_REPLY, redlot.Del)
	REG("EXISTS", INT_REPLY, redlot.Exists)
	REG("SETX", STATUS_REPLY, redlot.Setx)
	REG("SETEX", STATUS_REPLY, redlot.Setx) // Alias of SETX
	REG("TTL", INT_REPLY, redlot.Ttl)
	REG("EXPIRE", INT_REPLY, redlot.Expire)
	REGL("KEYS", LIST_REPLY, redlot.Keys)
	REGL("SCAN", LIST_REPLY, redlot.Scan)

}

func Serve(addr string, options *redlot.Options) {
	// Open LevelDB with options.
	redlot.Open(options)

	// Create sockets listener.
	l, err := net.Listen("tcp4", addr)
	if err != nil {
		log.Fatalf("Listen error: %v\n", err.Error())
	}
	defer l.Close()

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Fatalf("Wait for a connection error: %s\n", err.Error())
		}

		// Count connecion
		info.Lock()
		info.ConnCounter++
		info.Unlock()

		go func(c net.Conn) {
			for {
				req, err := newRequset(c)
				if err == io.EOF {
					break
				}
				if err != nil {
					continue
				}

				info.Lock()
				info.TotalCalls++
				info.Unlock()

				reply := RUN(req.Cmd, req.Args)
				reply.WriteTo(c)
			}

			c.Close()
			info.Lock()
			info.ConnCounter--
			info.Unlock()

		}(conn)

	}
}
<|endoftext|>


<commit_before>package resque

//
// Copyright 2017 Pedro Salgado
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

import (
	"fmt"
	"time"

	"gopkg.in/redis.v5"
)

const (
	// Queues name of the key for the set where Resque stores the currently available queues.
	Queues = "resque:queues"
)

// Queue Queue interface.
type Queue interface {
	Consumer() Consumer
	Name() string
	Producer() Producer
	Receive() (*Job, error)
	Send(job *Job) error
}

// RedisQueue a job RedisQueue.
type RedisQueue struct {
	redis   *redis.Client
	name    string
	timeout time.Duration
	// queueName name of the Queue.
	queueName string
	// Consumer Resque job consumer.
	consumer Consumer
	// Producer Resque job producer.
	producer Producer
}

// NewRedisQueue initializes a RedisQueue struct and updates the set of available Resque queues.
func NewRedisQueue(name string, rc *redis.Client, timeout time.Duration, c Consumer, p Producer) (*RedisQueue, error) {
	q := &RedisQueue{
		redis:     rc,
		name:      name,
		queueName: fmt.Sprintf("resque:Queue:%s", name),
		timeout:   timeout,
		consumer:  c,
		producer:  p,
	}
	if c != nil {
		c.SetQueue(q)
	}
	if p != nil {
		p.SetQueue(q)
	}

	exists, err := rc.SIsMember(Queues, name).Result()
	if err != nil {
		return nil, err
	}
	if !exists {
		_, err := rc.SAdd(Queues, name).Result()
		if err != nil {
			return nil, err
		}
	}

	return q, nil
}

// Consumer returns this Queue's consumer.
func (rq *RedisQueue) Consumer() Consumer {
	return rq.consumer
}

// Name returns the Queue name.
func (rq *RedisQueue) Name() string {
	return rq.name
}

// Producer returns this Queue's producer.
func (rq *RedisQueue) Producer() Producer {
	return rq.producer
}

// Peek returns from the RedisQueue the jobs at position(s) [start, stop].
func (rq *RedisQueue) Peek(start, stop int64) ([]Job, error) {

	var jobs []Job

	err := rq.redis.LRange(rq.queueName, start, stop).ScanSlice(jobs)
	if err != nil {
		return nil, err
	}

	return jobs, nil
}

// Receive gets a job from the RedisQueue.
func (rq *RedisQueue) Receive() (*Job, error) {

	var job Job

	v, err := rq.redis.BLPop(rq.timeout, rq.queueName).Result()
	if err != nil {
		if err.Error() == "redis: nil" {
			return nil, nil
		}
		return nil, err
	}

	err = job.UnmarshalBinary([]byte(v[1]))
	if err != nil {
		return nil, err
	}

	return &job, err
}

// Send places a job to the RedisQueue.
func (rq *RedisQueue) Send(job *Job) error {

	byteArr, err := job.MarshalBinary()
	if err != nil {
		return err
	}

	return rq.redis.RPush(rq.queueName, string(byteArr)).Err()
}

// Size returns the number of jobs in the RedisQueue.
func (rq *RedisQueue) Size() (int64, error) {
	v, err := rq.redis.LLen(rq.queueName).Result()
	if err != nil {
		return 0, err
	}

	return v, err
}
<commit_msg>queue.go: fixed expression to build queue name.<commit_after>package resque

//
// Copyright 2017 Pedro Salgado
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

import (
	"fmt"
	"time"

	"gopkg.in/redis.v5"
)

const (
	// Queues name of the key for the set where Resque stores the currently available queues.
	Queues = "resque:queues"
)

// Queue Queue interface.
type Queue interface {
	Consumer() Consumer
	Name() string
	Producer() Producer
	Receive() (*Job, error)
	Send(job *Job) error
}

// RedisQueue a job RedisQueue.
type RedisQueue struct {
	redis   *redis.Client
	name    string
	timeout time.Duration
	// queueName name of the Queue.
	queueName string
	// Consumer Resque job consumer.
	consumer Consumer
	// Producer Resque job producer.
	producer Producer
}

// NewRedisQueue initializes a RedisQueue struct and updates the set of available Resque queues.
func NewRedisQueue(name string, rc *redis.Client, timeout time.Duration, c Consumer, p Producer) (*RedisQueue, error) {
	q := &RedisQueue{
		redis:     rc,
		name:      name,
		queueName: fmt.Sprintf("resque:queue:%s", name),
		timeout:   timeout,
		consumer:  c,
		producer:  p,
	}
	if c != nil {
		c.SetQueue(q)
	}
	if p != nil {
		p.SetQueue(q)
	}

	exists, err := rc.SIsMember(Queues, name).Result()
	if err != nil {
		return nil, err
	}
	if !exists {
		_, err := rc.SAdd(Queues, name).Result()
		if err != nil {
			return nil, err
		}
	}

	return q, nil
}

// Consumer returns this Queue's consumer.
func (rq *RedisQueue) Consumer() Consumer {
	return rq.consumer
}

// Name returns the Queue name.
func (rq *RedisQueue) Name() string {
	return rq.name
}

// Producer returns this Queue's producer.
func (rq *RedisQueue) Producer() Producer {
	return rq.producer
}

// Peek returns from the RedisQueue the jobs at position(s) [start, stop].
func (rq *RedisQueue) Peek(start, stop int64) ([]Job, error) {

	var jobs []Job

	err := rq.redis.LRange(rq.queueName, start, stop).ScanSlice(jobs)
	if err != nil {
		return nil, err
	}

	return jobs, nil
}

// Receive gets a job from the RedisQueue.
func (rq *RedisQueue) Receive() (*Job, error) {

	var job Job

	v, err := rq.redis.BLPop(rq.timeout, rq.queueName).Result()
	if err != nil {
		if err.Error() == "redis: nil" {
			return nil, nil
		}
		return nil, err
	}

	err = job.UnmarshalBinary([]byte(v[1]))
	if err != nil {
		return nil, err
	}

	return &job, err
}

// Send places a job to the RedisQueue.
func (rq *RedisQueue) Send(job *Job) error {

	byteArr, err := job.MarshalBinary()
	if err != nil {
		return err
	}

	return rq.redis.RPush(rq.queueName, string(byteArr)).Err()
}

// Size returns the number of jobs in the RedisQueue.
func (rq *RedisQueue) Size() (int64, error) {
	v, err := rq.redis.LLen(rq.queueName).Result()
	if err != nil {
		return 0, err
	}

	return v, err
}
<|endoftext|>


<commit_before>package client

import (
	"net/http"
	"net/http/httputil"
	"time"

	"github.com/getlantern/balancer"
	"github.com/getlantern/detour"
	"github.com/getlantern/flashlight/proxy"
)

// getReverseProxy waits for a message from client.rpCh to arrive and then it
// writes it back to client.rpCh before returning it as a value. This way we
// always have a balancer at client.rpCh and, if we don't have one, it would
// block until one arrives.
func (client *Client) getReverseProxy() *httputil.ReverseProxy {
	rp := <-client.rpCh
	client.rpCh <- rp
	return rp
}

// initReverseProxy creates a reverse proxy that attempts to exit with any of
// the dialers provided by the balancer.
func (client *Client) initReverseProxy(bal *balancer.Balancer, dumpHeaders bool) {
	rp := &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			// do nothing
		},
		Transport: withDumpHeaders(
			dumpHeaders,
			&http.Transport{
				// We disable keepalives because some servers pretend to support
				// keep-alives but close their connections immediately, which
				// causes an error inside ReverseProxy.  This is not an issue
				// for HTTPS because  the browser is responsible for handling
				// the problem, which browsers like Chrome and Firefox already
				// know to do.
				//
				// See https://code.google.com/p/go/issues/detail?id=4677
				DisableKeepAlives: true,
				// TODO: would be good to make this sensitive to QOS, which
				// right now is only respected for HTTPS connections. The
				// challenge is that ReverseProxy reuses connections for
				// different requests, so we might have to configure different
				// ReverseProxies for different QOS's or something like that.
				Dial: detour.Dialer(bal.Dial),
			}),
		// Set a FlushInterval to prevent overly aggressive buffering of
		// responses, which helps keep memory usage down
		FlushInterval: 250 * time.Millisecond,
		ErrorLog:      log.AsStdLogger(),
	}

	if client.rpInitialized {
		log.Trace("Draining reverse proxy channel")
		<-client.rpCh
	} else {
		log.Trace("Creating reverse proxy channel")
		client.rpCh = make(chan *httputil.ReverseProxy, 1)
	}

	log.Trace("Publishing reverse proxy")

	client.rpCh <- rp

	// We don't need to protect client.rpInitialized from race conditions because
	// it's only accessed here in initReverseProxy, which always gets called
	// under Configure, which never gets called concurrently with itself.
	client.rpInitialized = true
}

// withDumpHeaders creates a RoundTripper that uses the supplied RoundTripper
// and that dumps headers is client is so configured.
func withDumpHeaders(shouldDumpHeaders bool, rt http.RoundTripper) http.RoundTripper {
	if !shouldDumpHeaders {
		return rt
	}
	return &headerDumpingRoundTripper{rt}
}

// headerDumpingRoundTripper is an http.RoundTripper that wraps another
// http.RoundTripper and dumps response headers to the log.
type headerDumpingRoundTripper struct {
	orig http.RoundTripper
}

func (rt *headerDumpingRoundTripper) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	proxy.DumpHeaders("Request", &req.Header)
	resp, err = rt.orig.RoundTrip(req)
	if err == nil {
		proxy.DumpHeaders("Response", &resp.Header)
	}
	return
}
<commit_msg>update reverseproxy to not use detour<commit_after>package client

import (
	"net/http"
	"net/http/httputil"
	"runtime"
	"time"

	"github.com/getlantern/balancer"
	"github.com/getlantern/detour"
	"github.com/getlantern/flashlight/proxy"
)

// getReverseProxy waits for a message from client.rpCh to arrive and then it
// writes it back to client.rpCh before returning it as a value. This way we
// always have a balancer at client.rpCh and, if we don't have one, it would
// block until one arrives.
func (client *Client) getReverseProxy() *httputil.ReverseProxy {
	rp := <-client.rpCh
	client.rpCh <- rp
	return rp
}

// initReverseProxy creates a reverse proxy that attempts to exit with any of
// the dialers provided by the balancer.
func (client *Client) initReverseProxy(bal *balancer.Balancer, dumpHeaders bool) {

	transport := &http.Transport{
		// We disable keepalives because some servers pretend to support
		// keep-alives but close their connections immediately, which
		// causes an error inside ReverseProxy.  This is not an issue
		// for HTTPS because  the browser is responsible for handling
		// the problem, which browsers like Chrome and Firefox already
		// know to do.
		//
		// See https://code.google.com/p/go/issues/detail?id=4677
		DisableKeepAlives: true,
	}

	// TODO: would be good to make this sensitive to QOS, which
	// right now is only respected for HTTPS connections. The
	// challenge is that ReverseProxy reuses connections for
	// different requests, so we might have to configure different
	// ReverseProxies for different QOS's or something like that.
	if runtime.GOOS == "android" {
		transport.Dial = bal.Dial
	} else {
		transport.Dial = detour.Dialer(bal.Dial)
	}

	rp := &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			// do nothing
		},
		Transport: withDumpHeaders(
			dumpHeaders,
			transport),
		// Set a FlushInterval to prevent overly aggressive buffering of
		// responses, which helps keep memory usage down
		FlushInterval: 250 * time.Millisecond,
		ErrorLog:      log.AsStdLogger(),
	}

	if client.rpInitialized {
		log.Trace("Draining reverse proxy channel")
		<-client.rpCh
	} else {
		log.Trace("Creating reverse proxy channel")
		client.rpCh = make(chan *httputil.ReverseProxy, 1)
	}

	log.Trace("Publishing reverse proxy")

	client.rpCh <- rp

	// We don't need to protect client.rpInitialized from race conditions because
	// it's only accessed here in initReverseProxy, which always gets called
	// under Configure, which never gets called concurrently with itself.
	client.rpInitialized = true
}

// withDumpHeaders creates a RoundTripper that uses the supplied RoundTripper
// and that dumps headers is client is so configured.
func withDumpHeaders(shouldDumpHeaders bool, rt http.RoundTripper) http.RoundTripper {
	if !shouldDumpHeaders {
		return rt
	}
	return &headerDumpingRoundTripper{rt}
}

// headerDumpingRoundTripper is an http.RoundTripper that wraps another
// http.RoundTripper and dumps response headers to the log.
type headerDumpingRoundTripper struct {
	orig http.RoundTripper
}

func (rt *headerDumpingRoundTripper) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	proxy.DumpHeaders("Request", &req.Header)
	resp, err = rt.orig.RoundTrip(req)
	if err == nil {
		proxy.DumpHeaders("Response", &resp.Header)
	}
	return
}
<|endoftext|>


<commit_before>package main

// This file is used for testing

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"

	yaml "gopkg.in/yaml.v2"

	"github.com/aryann/difflib"
	"golang.org/x/oauth2"

	githubApp "github.com/google/go-github/github"
)

type branches struct {
	branchList []string
	repo       *Repo
}

func getData() {
	// var conf *Setting
	// var configFile = "data/github/sairam/settings.yml"
	var accessTokenByUser = os.Getenv("GITHUB_USER_TOKEN") // this is temporary for validating responses

	ts := oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: accessTokenByUser},
	)
	tc := oauth2.NewClient(oauth2.NoContext, ts)

	client := githubApp.NewClient(tc)
	branchesURL := "https://api.github.com/repos/sairam/gitnotify/branches"
	v := new([]*BranchInfo)
	req, _ := http.NewRequest("GET", branchesURL, nil)
	client.Do(req, v)
	newBranches := make([]string, len(*v))
	for i, a := range *v {
		newBranches[i] = a.Name
	}

	var repo *Repo
	repo = &Repo{
		Repo: "sairam/gitnotify",
	}
	new := &branches{
		branchList: newBranches,
		repo:       repo,
	}

	old := &branches{
		repo: repo,
	}
	old.load()

	// check data difference with previously saved one
	diff := getNewStrings(old.branchList, new.branchList)
	fmt.Println(diff)

	new.save()

	// TODO
	// diff := diffData(v)
	// sendEmail(diff)
	// persistChanges(v)
}

func (b *branches) fileName() string {
	repo := b.repo
	fileName := strings.Replace(repo.Repo, "/", "__", 1)
	dir := fmt.Sprintf("%s/repo", "data/github/sairam")
	if _, err := os.Stat(dir); err != nil {
		os.Mkdir(dir, 0700)
	}
	return fmt.Sprintf("%s/%s.yml", dir, fileName)
}

func (b *branches) load() error {
	data, err := ioutil.ReadFile(b.fileName())
	if os.IsNotExist(err) {
		return err
	}

	err = yaml.Unmarshal(data, &b.branchList)
	if err != nil {
		return err
	}

	return nil
}

func (b *branches) save() error {
	out, err := yaml.Marshal(b.branchList)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(b.fileName(), out, 0600)
}

func getNewStrings(old, new []string) []string {
	var strs []string
	for _, s := range difflib.Diff(old, new) {
		if s.Delta == difflib.RightOnly {
			strs = append(strs, s.Payload)
		}
	}
	return strs
}

// func init() {
// 	getData()
// }
<commit_msg>refactor to modify data structure<commit_after>package main

// This file is used for testing

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"

	yaml "gopkg.in/yaml.v2"

	"github.com/aryann/difflib"
	"golang.org/x/oauth2"

	githubApp "github.com/google/go-github/github"
)

type branches struct {
	repo    *Repo
	auth    *Authentication
	client  *githubApp.Client
	option  string
	oldList []string
	newList []string
}

func getData() {
	// First get the authentication info
	auth := &Authentication{
		Provider: "github",
		UserName: "sairam",
		Token:    "",
	}

	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: auth.Token})
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	client := githubApp.NewClient(tc)

	// loop through repos and their branches
	repo := &Repo{
		Repo: "sairam/gitnotify",
	}

	branch := &branches{
		repo:   repo,
		client: client,
		auth:   auth,
	}

	fmt.Println(updateNewBranches(branch, "branches"))
	fmt.Println(updateNewBranches(branch, "tags"))

	// sendEmail(diff)
}

func updateNewBranches(branch *branches, option string) []string {
	branch.option = option
	branchesURL := "https://api.github.com/repos/sairam/gitnotify/" + option
	v := new([]*BranchInfo)
	req, _ := http.NewRequest("GET", branchesURL, nil)
	branch.client.Do(req, v)
	newBranches := make([]string, len(*v))
	for i, a := range *v {
		newBranches[i] = a.Name
	}

	branch.newList = newBranches
	branch.load()
	diff := branch.diff()
	branch.save()
	return diff
}

// check data difference with previously saved one
func (b *branches) diff() []string {
	return getNewStrings(b.oldList, b.newList)
}

func (b *branches) fileName() string {
	repo := b.repo
	fileName := strings.Replace(repo.Repo, "/", "__", 1)
	dir := fmt.Sprintf("data/%s/%s/repo", b.auth.Provider, b.auth.UserName)
	if _, err := os.Stat(dir); err != nil {
		os.Mkdir(dir, 0700)
	}
	return fmt.Sprintf("%s/%s-%s.yml", dir, fileName, b.option)
}

// load copies data into oldList
func (b *branches) load() error {
	data, err := ioutil.ReadFile(b.fileName())
	if os.IsNotExist(err) {
		return err
	}

	err = yaml.Unmarshal(data, &b.oldList)
	if err != nil {
		return err
	}

	return nil
}

// save copies data from newList into file
func (b *branches) save() error {
	out, err := yaml.Marshal(b.newList)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(b.fileName(), out, 0600)
}

func getNewStrings(old, new []string) []string {
	var strs []string
	for _, s := range difflib.Diff(old, new) {
		if s.Delta == difflib.RightOnly {
			strs = append(strs, s.Payload)
		}
	}
	return strs
}

func init() {
	getData()
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package s3

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	"github.com/keybase/release/version"
)

type Section struct {
	Header   string
	Releases []Release
}

type Release struct {
	Name    string
	URL     string
	Version string
	Date    string
	Commit  string
}

func NewClient() (client *s3.S3, err error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		return
	}
	client = s3.New(auth, aws.USEast)
	return
}

func WriteHTML(path string, bucketName string, prefixes string, suffix string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)
	if bucket == nil {
		return fmt.Errorf("Bucket %s not found", bucketName)
	}

	var sections []Section
	for _, prefix := range strings.Split(prefixes, ",") {
		resp, err := bucket.List(prefix, "", "", 0)
		if err != nil {
			return err
		}

		var releases []Release
		for _, k := range resp.Contents {
			if strings.HasSuffix(k.Key, suffix) {
				key := k.Key
				name := key[len(prefix):]
				urlString := fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
				version, date, commit, err := version.Parse(name)
				if err != nil {
					log.Printf("Couldn't get version from name: %s\n", name)
				}

				// Convert to Eastern
				locationNewYork, err := time.LoadLocation("America/New_York")
				if err != nil {
					log.Printf("Couldn't load location: %s", err)
				}
				date = date.In(locationNewYork)

				releases = append(releases,
					Release{
						Name:    name,
						URL:     urlString,
						Version: version,
						Date:    date.Format("Mon Jan _2 15:04:05 MST 2006"),
						Commit:  commit,
					})
			}
		}
		sections = append(sections, Section{
			Header:   prefix,
			Releases: reverseRelease(releases),
		})
	}

	return WriteHTMLForLinks(path, bucketName, sections)
}

var htmlTemplate = `
<!doctype html>

  <title>{{ .Title }}</title>

	{{ range $index, $sec := .Sections }}

{{ $sec.Header }}

		{{ range $index2, $rel := $sec.Releases }}

{{ $rel.Name }} {{ $rel.Version }} {{ $rel.Date }} {{ $rel.Commit }}

		{{ end }}

	{{ end }}

`

func WriteHTMLForLinks(path string, title string, sections []Section) error {
	vars := map[string]interface{}{
		"Title":    title,
		"Sections": sections,
	}

	t, err := template.New("t").Parse(htmlTemplate)
	if err != nil {
		return err
	}

	if path != "" {
		var data bytes.Buffer
		err = t.Execute(&data, vars)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(path, data.Bytes(), 0644)
	}
	return nil
}

type Link struct {
	Prefix string
	Suffix string
	Name   string
}

func CopyLatest(bucketName string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)

	linksForPrefix := []Link{
		Link{Prefix: "darwin/", Name: "Keybase.dmg"},
		Link{Prefix: "linux_binaries/deb/", Suffix: "_amd64.deb", Name: "keybase_amd64.deb"},
		Link{Prefix: "linux_binaries/rpm/", Suffix: ".x86_64.rpm", Name: "keybase_amd64.rpm"},
	}

	for _, link := range linksForPrefix {
		resp, err := bucket.List(link.Prefix, "", "", 0)
		if err != nil {
			return err
		}
		keys := reverseKey(resp.Contents)
		for _, k := range keys {
			if !strings.HasSuffix(k.Key, link.Suffix) {
				continue
			}

			url := urlString(k, bucketName, link.Prefix)
			// Instead of linking, we're making copies. S3 linking has some issues.
			// headers := map[string][]string{
			// 	"x-amz-website-redirect-location": []string{url},
			// }
			//err = bucket.PutHeader(name, []byte{}, headers, s3.PublicRead)
			log.Printf("Copying %s from %s (latest)\n", link.Name, k.Key)
			_, err = bucket.PutCopy(link.Name, s3.PublicRead, s3.CopyOptions{}, url)
			if err != nil {
				return err
			}
			break
		}
	}
	return nil
}

func urlString(k s3.Key, bucketName string, prefix string) string {
	key := k.Key
	name := key[len(prefix):]
	return fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
}

func reverseKey(a []s3.Key) []s3.Key {
	for left, right := 0, len(a)-1; left < right; left, right = left+1, right-1 {
		a[left], a[right] = a[right], a[left]
	}
	return a
}

func reverseRelease(a []Release) []Release {
	for left, right := 0, len(a)-1; left < right; left, right = left+1, right-1 {
		a[left], a[right] = a[right], a[left]
	}
	return a
}
<commit_msg>Limit s3 index to last 20 builds<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package s3

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	"github.com/keybase/release/version"
)

type Section struct {
	Header   string
	Releases []Release
}

type Release struct {
	Name    string
	URL     string
	Version string
	Date    string
	Commit  string
}

func NewClient() (client *s3.S3, err error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		return
	}
	client = s3.New(auth, aws.USEast)
	return
}

func WriteHTML(path string, bucketName string, prefixes string, suffix string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)
	if bucket == nil {
		return fmt.Errorf("Bucket %s not found", bucketName)
	}

	var sections []Section
	for _, prefix := range strings.Split(prefixes, ",") {
		resp, err := bucket.List(prefix, "", "", 0)
		if err != nil {
			return err
		}
		keys := reverseKey(resp.Contents, 20)

		var releases []Release
		for _, k := range keys {
			if strings.HasSuffix(k.Key, suffix) {
				key := k.Key
				name := key[len(prefix):]
				urlString := fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
				version, date, commit, err := version.Parse(name)
				if err != nil {
					log.Printf("Couldn't get version from name: %s\n", name)
				}

				// Convert to Eastern
				locationNewYork, err := time.LoadLocation("America/New_York")
				if err != nil {
					log.Printf("Couldn't load location: %s", err)
				}
				date = date.In(locationNewYork)

				releases = append(releases,
					Release{
						Name:    name,
						URL:     urlString,
						Version: version,
						Date:    date.Format("Mon Jan _2 15:04:05 MST 2006"),
						Commit:  commit,
					})
			}
		}
		sections = append(sections, Section{
			Header:   prefix,
			Releases: releases,
		})
	}

	return WriteHTMLForLinks(path, bucketName, sections)
}

var htmlTemplate = `
<!doctype html>

  <title>{{ .Title }}</title>

	{{ range $index, $sec := .Sections }}

{{ $sec.Header }}

		{{ range $index2, $rel := $sec.Releases }}

{{ $rel.Name }} {{ $rel.Version }} {{ $rel.Date }} {{ $rel.Commit }}

		{{ end }}

	{{ end }}

`

func WriteHTMLForLinks(path string, title string, sections []Section) error {
	vars := map[string]interface{}{
		"Title":    title,
		"Sections": sections,
	}

	t, err := template.New("t").Parse(htmlTemplate)
	if err != nil {
		return err
	}

	if path != "" {
		var data bytes.Buffer
		err = t.Execute(&data, vars)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(path, data.Bytes(), 0644)
	}
	return nil
}

type Link struct {
	Prefix string
	Suffix string
	Name   string
}

func CopyLatest(bucketName string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)

	linksForPrefix := []Link{
		Link{Prefix: "darwin/", Name: "Keybase.dmg"},
		Link{Prefix: "linux_binaries/deb/", Suffix: "_amd64.deb", Name: "keybase_amd64.deb"},
		Link{Prefix: "linux_binaries/rpm/", Suffix: ".x86_64.rpm", Name: "keybase_amd64.rpm"},
	}

	for _, link := range linksForPrefix {
		resp, err := bucket.List(link.Prefix, "", "", 0)
		if err != nil {
			return err
		}
		keys := reverseKey(resp.Contents, 0)
		for _, k := range keys {
			if !strings.HasSuffix(k.Key, link.Suffix) {
				continue
			}

			url := urlString(k, bucketName, link.Prefix)
			// Instead of linking, we're making copies. S3 linking has some issues.
			// headers := map[string][]string{
			// 	"x-amz-website-redirect-location": []string{url},
			// }
			//err = bucket.PutHeader(name, []byte{}, headers, s3.PublicRead)
			log.Printf("Copying %s from %s (latest)\n", link.Name, k.Key)
			_, err = bucket.PutCopy(link.Name, s3.PublicRead, s3.CopyOptions{}, url)
			if err != nil {
				return err
			}
			break
		}
	}
	return nil
}

func urlString(k s3.Key, bucketName string, prefix string) string {
	key := k.Key
	name := key[len(prefix):]
	return fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
}

func reverseKey(a []s3.Key, truncate int) []s3.Key {
	for left, right := 0, len(a)-1; left < right; left, right = left+1, right-1 {
		a[left], a[right] = a[right], a[left]
	}
	if truncate > 0 && len(a) > truncate {
		a = a[0:truncate]
	}
	return a
}
<|endoftext|>


<commit_before>package main

import (
    "log"
    "net"
    "time"
    "sync"
    "container/list"
    "huabot-sched/db"
)

type Sched struct {
    worker_count int
    timer *time.Timer
    queue *list.List
    jobQueue *list.List
    sockFile string
    locker   *sync.Mutex
}

func NewSched(sockFile string) *Sched {
    sched = new(Sched)
    sched.worker_count = 0
    sched.timer = time.NewTimer(1 * time.Hour)
    sched.queue = list.New()
    sched.jobQueue = list.New()
    sched.sockFile = sockFile
    sched.locker = new(sync.Mutex)
    return sched
}

func (sched *Sched) Serve() {
    sockCheck(sched.sockFile)
    sched.checkJobQueue()
    go sched.handle()
    listen, err := net.Listen("unix", sched.sockFile)
    if err != nil {
        log.Fatal(err)
    }
    defer listen.Close()
    log.Printf("huabot-sched started on %s\n", sched.sockFile)
    for {
        conn, err := listen.Accept()
        if err != nil {
            log.Fatal(err)
        }
        sched.HandleConnection(conn)
    }
}

func (sched *Sched) Notify() {
    sched.timer.Reset(time.Millisecond)
}

func (sched *Sched) DieWorker(worker *Worker) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.worker_count -= 1
    log.Printf("worker_count: %d\n", sched.worker_count)
    sched.removeQueue(worker)
    worker.Close()
}

func (sched *Sched) HandleConnection(conn net.Conn) {
    worker := NewWorker(sched, Conn{Conn: conn})
    sched.worker_count += 1
    log.Printf("worker_count: %d\n", sched.worker_count)
    go worker.Handle()
}

func (sched *Sched) Done(jobId int) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.locker.Lock()
    removeListJob(sched.jobQueue, jobId)
    db.DelJob(jobId)
    return
}

func (sched *Sched) isDoJob(job db.Job) bool {
    now := time.Now()
    current := int(now.Unix())
    ret := false
    for e := sched.jobQueue.Front(); e != nil; e = e.Next() {
        chk := e.Value.(db.Job)
        if chk.Timeout > 0 && chk.SchedAt + chk.Timeout > current {
            newJob, _ := db.GetJob(chk.Id)
            if newJob.Status == "doing" {
                newJob.Status = "ready"
                newJob.Save()
            }
            sched.jobQueue.Remove(e)
            continue
        }
        if chk.Id == job.Id {
            old := e.Value.(db.Job)
            if old.Timeout > 0 && old.SchedAt + old.Timeout < int(now.Unix()) {
                ret = false
            } else {
                ret = true
            }
        }
    }
    return ret
}

func (sched *Sched) SubmitJob(worker *Worker, job db.Job) {
    defer sched.locker.Unlock()
    sched.locker.Lock()
    if job.Name == "" {
        job.Delete()
        return
    }
    if sched.isDoJob(job) {
        return
    }
    if !worker.alive {
        return
    }
    if err := worker.HandleDo(job); err != nil {
        worker.alive = false
        sched.DieWorker(worker)
        return
    }
    job.Status = "doing"
    job.Save()
    sched.jobQueue.PushBack(job)
    sched.removeQueue(worker)
}

func (sched *Sched) handle() {
    var current time.Time
    var timestamp int
    for {
        for e := sched.queue.Front(); e != nil; e = e.Next() {
            worker := e.Value.(*Worker)
            jobs, err := db.RangeSchedJob("ready", 0, 0)
            if err != nil || len(jobs) == 0 {
                sched.timer.Reset(time.Minute)
                current =<-sched.timer.C
                continue
            }
            timestamp = int(time.Now().Unix())
            if jobs[0].SchedAt < timestamp {
                sched.SubmitJob(worker, jobs[0])
            } else {
                sched.timer.Reset(time.Second * time.Duration(jobs[0].SchedAt - timestamp))
                current =<-sched.timer.C
                timestamp = int(current.Unix())
                if jobs[0].SchedAt <= timestamp {
                    sched.SubmitJob(worker, jobs[0])
                }
            }
        }
        if sched.queue.Len() == 0 {
            sched.timer.Reset(time.Minute)
            current =<-sched.timer.C
        }
    }
}

func (sched *Sched) Fail(jobId int) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.locker.Lock()
    removeListJob(sched.jobQueue, jobId)
    job, _ := db.GetJob(jobId)
    job.Status = "ready"
    job.Save()
    return
}

func (sched *Sched) SchedLater(jobId int, delay int) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.locker.Lock()
    removeListJob(sched.jobQueue, jobId)
    job, _ := db.GetJob(jobId)
    job.Status = "ready"
    var now = time.Now()
    job.SchedAt = int(now.Unix()) + delay
    job.Save()
    return
}

func (sched *Sched) removeQueue(worker *Worker) {
    for e := sched.queue.Front(); e != nil; e = e.Next() {
        if e.Value.(*Worker) == worker {
            sched.queue.Remove(e)
        }
    }
}

func (sched *Sched) checkJobQueue() {
    start := 0
    limit := 20
    total, _ := db.CountSchedJob("doing")
    updateQueue := make([]db.Job, 0)
    removeQueue := make([]db.Job, 0)
    var now = time.Now()
    current := int(now.Unix())

    for start = 0; start < total; start += limit {
        jobs, _ := db.RangeSchedJob("doing", start, start + limit)
        for _, job := range jobs {
            if job.Name == "" {
                removeQueue = append(removeQueue, job)
                continue
            }
            if job.SchedAt + job.Timeout < current {
                updateQueue = append(updateQueue, job)
            } else {
                sched.jobQueue.PushBack(job)
            }
        }
    }

    for _, job := range updateQueue {
        job.Status = "ready"
        job.Save()
    }

    for _, job := range removeQueue {
        job.Delete()
    }
}

func (sched *Sched) Close() {
}
<commit_msg>rename worker_count -> workerCount<commit_after>package main

import (
    "log"
    "net"
    "time"
    "sync"
    "container/list"
    "huabot-sched/db"
)

type Sched struct {
    workerCount int
    timer *time.Timer
    queue *list.List
    jobQueue *list.List
    sockFile string
    locker   *sync.Mutex
}

func NewSched(sockFile string) *Sched {
    sched = new(Sched)
    sched.workerCount = 0
    sched.timer = time.NewTimer(1 * time.Hour)
    sched.queue = list.New()
    sched.jobQueue = list.New()
    sched.sockFile = sockFile
    sched.locker = new(sync.Mutex)
    return sched
}

func (sched *Sched) Serve() {
    sockCheck(sched.sockFile)
    sched.checkJobQueue()
    go sched.handle()
    listen, err := net.Listen("unix", sched.sockFile)
    if err != nil {
        log.Fatal(err)
    }
    defer listen.Close()
    log.Printf("huabot-sched started on %s\n", sched.sockFile)
    for {
        conn, err := listen.Accept()
        if err != nil {
            log.Fatal(err)
        }
        sched.HandleConnection(conn)
    }
}

func (sched *Sched) Notify() {
    sched.timer.Reset(time.Millisecond)
}

func (sched *Sched) DieWorker(worker *Worker) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.workerCount -= 1
    log.Printf("Total worker: %d\n", sched.workerCount)
    sched.removeQueue(worker)
    worker.Close()
}

func (sched *Sched) HandleConnection(conn net.Conn) {
    worker := NewWorker(sched, Conn{Conn: conn})
    sched.workerCount += 1
    log.Printf("Total worker: %d\n", sched.workerCount)
    go worker.Handle()
}

func (sched *Sched) Done(jobId int) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.locker.Lock()
    removeListJob(sched.jobQueue, jobId)
    db.DelJob(jobId)
    return
}

func (sched *Sched) isDoJob(job db.Job) bool {
    now := time.Now()
    current := int(now.Unix())
    ret := false
    for e := sched.jobQueue.Front(); e != nil; e = e.Next() {
        chk := e.Value.(db.Job)
        if chk.Timeout > 0 && chk.SchedAt + chk.Timeout > current {
            newJob, _ := db.GetJob(chk.Id)
            if newJob.Status == "doing" {
                newJob.Status = "ready"
                newJob.Save()
            }
            sched.jobQueue.Remove(e)
            continue
        }
        if chk.Id == job.Id {
            old := e.Value.(db.Job)
            if old.Timeout > 0 && old.SchedAt + old.Timeout < int(now.Unix()) {
                ret = false
            } else {
                ret = true
            }
        }
    }
    return ret
}

func (sched *Sched) SubmitJob(worker *Worker, job db.Job) {
    defer sched.locker.Unlock()
    sched.locker.Lock()
    if job.Name == "" {
        job.Delete()
        return
    }
    if sched.isDoJob(job) {
        return
    }
    if !worker.alive {
        return
    }
    if err := worker.HandleDo(job); err != nil {
        worker.alive = false
        sched.DieWorker(worker)
        return
    }
    job.Status = "doing"
    job.Save()
    sched.jobQueue.PushBack(job)
    sched.removeQueue(worker)
}

func (sched *Sched) handle() {
    var current time.Time
    var timestamp int
    for {
        for e := sched.queue.Front(); e != nil; e = e.Next() {
            worker := e.Value.(*Worker)
            jobs, err := db.RangeSchedJob("ready", 0, 0)
            if err != nil || len(jobs) == 0 {
                sched.timer.Reset(time.Minute)
                current =<-sched.timer.C
                continue
            }
            timestamp = int(time.Now().Unix())
            if jobs[0].SchedAt < timestamp {
                sched.SubmitJob(worker, jobs[0])
            } else {
                sched.timer.Reset(time.Second * time.Duration(jobs[0].SchedAt - timestamp))
                current =<-sched.timer.C
                timestamp = int(current.Unix())
                if jobs[0].SchedAt <= timestamp {
                    sched.SubmitJob(worker, jobs[0])
                }
            }
        }
        if sched.queue.Len() == 0 {
            sched.timer.Reset(time.Minute)
            current =<-sched.timer.C
        }
    }
}

func (sched *Sched) Fail(jobId int) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.locker.Lock()
    removeListJob(sched.jobQueue, jobId)
    job, _ := db.GetJob(jobId)
    job.Status = "ready"
    job.Save()
    return
}

func (sched *Sched) SchedLater(jobId int, delay int) {
    defer sched.Notify()
    defer sched.locker.Unlock()
    sched.locker.Lock()
    removeListJob(sched.jobQueue, jobId)
    job, _ := db.GetJob(jobId)
    job.Status = "ready"
    var now = time.Now()
    job.SchedAt = int(now.Unix()) + delay
    job.Save()
    return
}

func (sched *Sched) removeQueue(worker *Worker) {
    for e := sched.queue.Front(); e != nil; e = e.Next() {
        if e.Value.(*Worker) == worker {
            sched.queue.Remove(e)
        }
    }
}

func (sched *Sched) checkJobQueue() {
    start := 0
    limit := 20
    total, _ := db.CountSchedJob("doing")
    updateQueue := make([]db.Job, 0)
    removeQueue := make([]db.Job, 0)
    var now = time.Now()
    current := int(now.Unix())

    for start = 0; start < total; start += limit {
        jobs, _ := db.RangeSchedJob("doing", start, start + limit)
        for _, job := range jobs {
            if job.Name == "" {
                removeQueue = append(removeQueue, job)
                continue
            }
            if job.SchedAt + job.Timeout < current {
                updateQueue = append(updateQueue, job)
            } else {
                sched.jobQueue.PushBack(job)
            }
        }
    }

    for _, job := range updateQueue {
        job.Status = "ready"
        job.Save()
    }

    for _, job := range removeQueue {
        job.Delete()
    }
}

func (sched *Sched) Close() {
}
<|endoftext|>


<commit_before>package periodic

import (
    "log"
    "net"
    "time"
    "sync"
    "strings"
    "container/heap"
    "github.com/Lupino/periodic/driver"
    "github.com/Lupino/periodic/protocol"
)

type Sched struct {
    jobTimer   *time.Timer
    grabQueue  *GrabQueue
    procQueue  map[int64]driver.Job
    revertPQ   PriorityQueue
    revTimer   *time.Timer
    entryPoint string
    JobLocker  *sync.Mutex
    TimerLocker *sync.Mutex
    stats      map[string]*FuncStat
    FuncLocker *sync.Mutex
    driver     driver.StoreDriver
    jobPQ      map[string]*PriorityQueue
    PQLocker   *sync.Mutex
    timeout    time.Duration
    alive      bool
    cacheItem  *Item
}

func NewSched(entryPoint string, store driver.StoreDriver, timeout time.Duration) *Sched {
    sched := new(Sched)
    sched.jobTimer = time.NewTimer(1 * time.Hour)
    sched.revTimer = time.NewTimer(1 * time.Hour)
    sched.grabQueue = NewGrabQueue()
    sched.procQueue = make(map[int64]driver.Job)
    sched.revertPQ = make(PriorityQueue, 0)
    heap.Init(&sched.revertPQ)
    sched.entryPoint = entryPoint
    sched.JobLocker = new(sync.Mutex)
    sched.PQLocker = new(sync.Mutex)
    sched.FuncLocker = new(sync.Mutex)
    sched.TimerLocker = new(sync.Mutex)
    sched.stats = make(map[string]*FuncStat)
    sched.driver = store
    sched.jobPQ = make(map[string]*PriorityQueue)
    sched.timeout = timeout
    sched.alive = true
    sched.cacheItem = nil
    return sched
}

func (sched *Sched) Serve() {
    parts := strings.SplitN(sched.entryPoint, "://", 2)
    if parts[0] == "unix" {
        sockCheck(parts[1])
    }
    sched.loadJobQueue()
    go sched.handleJobPQ()
    go sched.handleRevertPQ()
    listen, err := net.Listen(parts[0], parts[1])
    if err != nil {
        log.Fatal(err)
    }
    defer listen.Close()
    log.Printf("Periodic task system started on %s\n", sched.entryPoint)
    for {
        if !sched.alive {
            break
        }
        conn, err := listen.Accept()
        if err != nil {
            log.Fatal(err)
        }
        if sched.timeout > 0 {
            conn.SetDeadline(time.Now().Add(sched.timeout * time.Second))
        }
        sched.HandleConnection(conn)
    }
}

func (sched *Sched) NotifyJobTimer() {
    sched.resetJobTimer(time.Millisecond)
}

func (sched *Sched) resetJobTimer(d time.Duration) {
    defer sched.TimerLocker.Unlock()
    sched.TimerLocker.Lock()
    sched.jobTimer.Reset(d)
}

func (sched *Sched) NotifyRevertTimer() {
    sched.resetRevertTimer(time.Millisecond)
}

func (sched *Sched) resetRevertTimer(d time.Duration) {
    defer sched.TimerLocker.Unlock()
    sched.TimerLocker.Lock()
    sched.revTimer.Reset(d)
}

func (sched *Sched) HandleConnection(conn net.Conn) {
    c := Conn{Conn: conn}
    payload, err := c.Receive()
    if err != nil {
        return
    }
    switch protocol.ClientType(payload[0]) {
    case protocol.TYPE_CLIENT:
        client := NewClient(sched, c)
        go client.Handle()
        break
    case protocol.TYPE_WORKER:
        worker := NewWorker(sched, c)
        go worker.Handle()
        break
    default:
        log.Printf("Unsupport client %d\n", payload[0])
        c.Close()
        break
    }
}

func (sched *Sched) Done(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if _, ok := sched.procQueue[jobId]; ok {
        delete(sched.procQueue, jobId)
    }
    job, err := sched.driver.Get(jobId)
    if err == nil {
        sched.driver.Delete(jobId)
        sched.DecrStatJob(job)
        sched.DecrStatProc(job)
        sched.removeRevertPQ(job)
    }
    return
}

func (sched *Sched) SubmitJob(grabItem GrabItem, job driver.Job) bool {
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if job.Name == "" {
        sched.driver.Delete(job.Id)
        return true
    }
    if _, ok := sched.procQueue[job.Id]; ok {
        return true
    }

    if !grabItem.w.alive {
        return false
    }
    if err := grabItem.w.HandleDo(grabItem.msgId, job); err != nil {
        grabItem.w.alive = false
        return false
    }
    now := time.Now()
    current := int64(now.Unix())
    job.Status = driver.JOB_STATUS_PROC
    job.RunAt = current
    sched.driver.Save(&job)
    sched.IncrStatProc(job)
    sched.pushRevertPQ(job)
    sched.NotifyRevertTimer()
    sched.procQueue[job.Id] = job
    sched.grabQueue.Remove(grabItem)
    return true
}

func (sched *Sched) clearCacheItem() {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    sched.cacheItem = nil
}

func (sched *Sched) lessItem() (lessItem *Item) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if sched.cacheItem != nil {
        return sched.cacheItem
    }
    maybeItem := make(map[string]*Item)
    for Func, stat := range sched.stats {
        if stat.Worker.Int() == 0 {
            continue
        }
        pq, ok := sched.jobPQ[Func]
        if !ok || pq.Len() == 0 {
            continue
        }

        item := heap.Pop(pq).(*Item)

        maybeItem[Func] = item

    }

    if len(maybeItem) == 0 {
        return nil
    }

    var lessFunc string

    for Func, item := range maybeItem {
        if lessItem == nil {
            lessItem = item
            lessFunc = Func
            continue
        }
        if lessItem.priority > item.priority {
            lessItem = item
            lessFunc = Func
        }
    }

    for Func, item := range maybeItem {
        if Func == lessFunc {
            continue
        }
        pq := sched.jobPQ[Func]
        heap.Push(pq, item)
    }
    sched.cacheItem = lessItem
    return
}

func (sched *Sched) handleJobPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        if sched.grabQueue.Len() == 0 {
            sched.resetJobTimer(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        lessItem := sched.lessItem()

        if lessItem == nil {
            sched.resetJobTimer(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        schedJob, err := sched.driver.Get(lessItem.value)

        if err != nil {
            sched.clearCacheItem()
            log.Printf("handleJobPQ error job: %d %v\n", lessItem.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if schedJob.SchedAt > timestamp {
            sched.resetJobTimer(time.Second * time.Duration(schedJob.SchedAt - timestamp))
            current =<-sched.jobTimer.C
            timestamp = int64(current.Unix())
            if schedJob.SchedAt > timestamp {
                sched.pushJobPQ(schedJob)
                continue
            }
        }

        grabItem, err := sched.grabQueue.Get(schedJob.Func)
        if err == nil {
            if sched.SubmitJob(grabItem, schedJob) {
                sched.clearCacheItem()
            } else {
                sched.pushJobPQ(schedJob)
            }
        } else {
            sched.pushJobPQ(schedJob)
        }
    }
}

func (sched *Sched) handleRevertPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        sched.PQLocker.Lock()
        pqLen := sched.revertPQ.Len()
        sched.PQLocker.Unlock()
        if pqLen == 0 {
            sched.resetRevertTimer(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        sched.PQLocker.Lock()
        item := heap.Pop(&sched.revertPQ).(*Item)
        sched.PQLocker.Unlock()

        if item == nil {
            sched.resetRevertTimer(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        revertJob, err := sched.driver.Get(item.value)

        if err != nil {
            log.Printf("handleRevertPQ error: job: %d %v\n", item.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if item.priority > timestamp {
            sched.resetRevertTimer(time.Second * time.Duration(item.priority - timestamp))
            current =<-sched.revTimer.C
            timestamp = int64(current.Unix())
            if item.priority > timestamp {
                sched.pushRevertPQ(revertJob)
                continue
            }
        }

        sched.DecrStatProc(revertJob)
        revertJob.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&revertJob)
        sched.pushJobPQ(revertJob)
        if _, ok := sched.procQueue[revertJob.Id]; ok {
            delete(sched.procQueue, revertJob.Id)
        }
    }
}

func (sched *Sched) Fail(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if _, ok := sched.procQueue[jobId]; ok {
        delete(sched.procQueue, jobId)
    }
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) getFuncStat(Func string) *FuncStat {
    defer sched.FuncLocker.Unlock()
    sched.FuncLocker.Lock()
    stat, ok := sched.stats[Func]
    if !ok {
        stat = NewFuncStat(Func)
        sched.stats[Func] = stat
    }
    return stat
}

func (sched *Sched) IncrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Incr()
}

func (sched *Sched) DecrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Decr()
}

func (sched *Sched) IncrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Incr()
}

func (sched *Sched) DecrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Decr()
}

func (sched *Sched) IncrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Incr()
    }
}

func (sched *Sched) DecrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Decr()
    }
}

func (sched *Sched) SchedLater(jobId int64, delay int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if _, ok := sched.procQueue[jobId]; ok {
        delete(sched.procQueue, jobId)
    }
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    var now = time.Now()
    job.SchedAt = int64(now.Unix()) + delay
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) pushJobPQ(job driver.Job) bool {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_READY {
        item := &Item{
            value: job.Id,
            priority: job.SchedAt,
        }
        if sched.cacheItem != nil && item.priority < sched.cacheItem.priority {
            if job.Id == sched.cacheItem.value {
                return true
            }
            job, _ = sched.driver.Get(sched.cacheItem.value)
            sched.cacheItem = item
            if job.Id <= 0 || job.Status != driver.JOB_STATUS_READY {
                return false
            }
        }
        pq, ok := sched.jobPQ[job.Func]
        if !ok {
            pq1 := make(PriorityQueue, 0)
            pq = &pq1
            sched.jobPQ[job.Func] = pq
            heap.Init(pq)
        }
        heap.Push(pq, item)
        return true
    }
    return false
}

func (sched *Sched) pushRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        runAt := job.RunAt
        if runAt == 0 {
            runAt = job.SchedAt
        }
        item := &Item{
            value: job.Id,
            priority: runAt + job.Timeout,
        }
        heap.Push(&sched.revertPQ, item)
    }
}

func (sched *Sched) removeRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        for _, item := range sched.revertPQ {
            if item.value == job.Id {
                heap.Remove(&sched.revertPQ, item.index)
                break
            }
        }
    }
}

func (sched *Sched) loadJobQueue() {
    updateQueue := make([]driver.Job, 0)
    removeQueue := make([]driver.Job, 0)
    var now = time.Now()
    current := int64(now.Unix())

    iter := sched.driver.NewIterator(nil)
    for {
        if !iter.Next() {
            break
        }
        job := iter.Value()
        if job.Name == "" {
            removeQueue = append(removeQueue, job)
            continue
        }
        sched.IncrStatJob(job)
        sched.pushJobPQ(job)
        runAt := job.RunAt
        if runAt < job.SchedAt {
            runAt = job.SchedAt
        }
        if runAt + job.Timeout < current {
            updateQueue = append(updateQueue, job)
        } else {
            sched.procQueue[job.Id] = job
            sched.IncrStatProc(job)
            sched.pushRevertPQ(job)
        }
    }

    iter.Close()

    for _, job := range updateQueue {
        job.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&job)
    }

    for _, job := range removeQueue {
        sched.driver.Delete(job.Id)
    }
}

func (sched *Sched) Close() {
    sched.alive = false
    sched.driver.Close()
    log.Printf("Periodic task system shutdown\n")
}
<commit_msg>Fix. procQueue DATA RACE<commit_after>package periodic

import (
    "log"
    "net"
    "time"
    "sync"
    "strings"
    "container/heap"
    "github.com/Lupino/periodic/driver"
    "github.com/Lupino/periodic/protocol"
)

type Sched struct {
    jobTimer   *time.Timer
    grabQueue  *GrabQueue
    procQueue  map[int64]driver.Job
    revertPQ   PriorityQueue
    revTimer   *time.Timer
    entryPoint string
    JobLocker  *sync.Mutex
    TimerLocker *sync.Mutex
    stats      map[string]*FuncStat
    FuncLocker *sync.Mutex
    driver     driver.StoreDriver
    jobPQ      map[string]*PriorityQueue
    PQLocker   *sync.Mutex
    timeout    time.Duration
    alive      bool
    cacheItem  *Item
}

func NewSched(entryPoint string, store driver.StoreDriver, timeout time.Duration) *Sched {
    sched := new(Sched)
    sched.jobTimer = time.NewTimer(1 * time.Hour)
    sched.revTimer = time.NewTimer(1 * time.Hour)
    sched.grabQueue = NewGrabQueue()
    sched.procQueue = make(map[int64]driver.Job)
    sched.revertPQ = make(PriorityQueue, 0)
    heap.Init(&sched.revertPQ)
    sched.entryPoint = entryPoint
    sched.JobLocker = new(sync.Mutex)
    sched.PQLocker = new(sync.Mutex)
    sched.FuncLocker = new(sync.Mutex)
    sched.TimerLocker = new(sync.Mutex)
    sched.stats = make(map[string]*FuncStat)
    sched.driver = store
    sched.jobPQ = make(map[string]*PriorityQueue)
    sched.timeout = timeout
    sched.alive = true
    sched.cacheItem = nil
    return sched
}

func (sched *Sched) Serve() {
    parts := strings.SplitN(sched.entryPoint, "://", 2)
    if parts[0] == "unix" {
        sockCheck(parts[1])
    }
    sched.loadJobQueue()
    go sched.handleJobPQ()
    go sched.handleRevertPQ()
    listen, err := net.Listen(parts[0], parts[1])
    if err != nil {
        log.Fatal(err)
    }
    defer listen.Close()
    log.Printf("Periodic task system started on %s\n", sched.entryPoint)
    for {
        if !sched.alive {
            break
        }
        conn, err := listen.Accept()
        if err != nil {
            log.Fatal(err)
        }
        if sched.timeout > 0 {
            conn.SetDeadline(time.Now().Add(sched.timeout * time.Second))
        }
        sched.HandleConnection(conn)
    }
}

func (sched *Sched) NotifyJobTimer() {
    sched.resetJobTimer(time.Millisecond)
}

func (sched *Sched) resetJobTimer(d time.Duration) {
    defer sched.TimerLocker.Unlock()
    sched.TimerLocker.Lock()
    sched.jobTimer.Reset(d)
}

func (sched *Sched) NotifyRevertTimer() {
    sched.resetRevertTimer(time.Millisecond)
}

func (sched *Sched) resetRevertTimer(d time.Duration) {
    defer sched.TimerLocker.Unlock()
    sched.TimerLocker.Lock()
    sched.revTimer.Reset(d)
}

func (sched *Sched) HandleConnection(conn net.Conn) {
    c := Conn{Conn: conn}
    payload, err := c.Receive()
    if err != nil {
        return
    }
    switch protocol.ClientType(payload[0]) {
    case protocol.TYPE_CLIENT:
        client := NewClient(sched, c)
        go client.Handle()
        break
    case protocol.TYPE_WORKER:
        worker := NewWorker(sched, c)
        go worker.Handle()
        break
    default:
        log.Printf("Unsupport client %d\n", payload[0])
        c.Close()
        break
    }
}

func (sched *Sched) Done(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if _, ok := sched.procQueue[jobId]; ok {
        delete(sched.procQueue, jobId)
    }
    job, err := sched.driver.Get(jobId)
    if err == nil {
        sched.driver.Delete(jobId)
        sched.DecrStatJob(job)
        sched.DecrStatProc(job)
        sched.removeRevertPQ(job)
    }
    return
}

func (sched *Sched) SubmitJob(grabItem GrabItem, job driver.Job) bool {
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if job.Name == "" {
        sched.driver.Delete(job.Id)
        return true
    }
    if _, ok := sched.procQueue[job.Id]; ok {
        return true
    }

    if !grabItem.w.alive {
        return false
    }
    if err := grabItem.w.HandleDo(grabItem.msgId, job); err != nil {
        grabItem.w.alive = false
        return false
    }
    now := time.Now()
    current := int64(now.Unix())
    job.Status = driver.JOB_STATUS_PROC
    job.RunAt = current
    sched.driver.Save(&job)
    sched.IncrStatProc(job)
    sched.pushRevertPQ(job)
    sched.NotifyRevertTimer()
    sched.procQueue[job.Id] = job
    sched.grabQueue.Remove(grabItem)
    return true
}

func (sched *Sched) clearCacheItem() {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    sched.cacheItem = nil
}

func (sched *Sched) lessItem() (lessItem *Item) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if sched.cacheItem != nil {
        return sched.cacheItem
    }
    maybeItem := make(map[string]*Item)
    for Func, stat := range sched.stats {
        if stat.Worker.Int() == 0 {
            continue
        }
        pq, ok := sched.jobPQ[Func]
        if !ok || pq.Len() == 0 {
            continue
        }

        item := heap.Pop(pq).(*Item)

        maybeItem[Func] = item

    }

    if len(maybeItem) == 0 {
        return nil
    }

    var lessFunc string

    for Func, item := range maybeItem {
        if lessItem == nil {
            lessItem = item
            lessFunc = Func
            continue
        }
        if lessItem.priority > item.priority {
            lessItem = item
            lessFunc = Func
        }
    }

    for Func, item := range maybeItem {
        if Func == lessFunc {
            continue
        }
        pq := sched.jobPQ[Func]
        heap.Push(pq, item)
    }
    sched.cacheItem = lessItem
    return
}

func (sched *Sched) handleJobPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        if sched.grabQueue.Len() == 0 {
            sched.resetJobTimer(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        lessItem := sched.lessItem()

        if lessItem == nil {
            sched.resetJobTimer(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        schedJob, err := sched.driver.Get(lessItem.value)

        if err != nil {
            sched.clearCacheItem()
            log.Printf("handleJobPQ error job: %d %v\n", lessItem.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if schedJob.SchedAt > timestamp {
            sched.resetJobTimer(time.Second * time.Duration(schedJob.SchedAt - timestamp))
            current =<-sched.jobTimer.C
            timestamp = int64(current.Unix())
            if schedJob.SchedAt > timestamp {
                sched.pushJobPQ(schedJob)
                continue
            }
        }

        grabItem, err := sched.grabQueue.Get(schedJob.Func)
        if err == nil {
            if sched.SubmitJob(grabItem, schedJob) {
                sched.clearCacheItem()
            } else {
                sched.pushJobPQ(schedJob)
            }
        } else {
            sched.pushJobPQ(schedJob)
        }
    }
}

func (sched *Sched) handleRevertPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        sched.PQLocker.Lock()
        pqLen := sched.revertPQ.Len()
        sched.PQLocker.Unlock()
        if pqLen == 0 {
            sched.resetRevertTimer(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        sched.PQLocker.Lock()
        item := heap.Pop(&sched.revertPQ).(*Item)
        sched.PQLocker.Unlock()

        if item == nil {
            sched.resetRevertTimer(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        revertJob, err := sched.driver.Get(item.value)

        if err != nil {
            log.Printf("handleRevertPQ error: job: %d %v\n", item.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if item.priority > timestamp {
            sched.resetRevertTimer(time.Second * time.Duration(item.priority - timestamp))
            current =<-sched.revTimer.C
            timestamp = int64(current.Unix())
            if item.priority > timestamp {
                sched.pushRevertPQ(revertJob)
                continue
            }
        }

        sched.DecrStatProc(revertJob)
        revertJob.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&revertJob)
        sched.pushJobPQ(revertJob)
        sched.JobLocker.Lock()
        if _, ok := sched.procQueue[revertJob.Id]; ok {
            delete(sched.procQueue, revertJob.Id)
        }
        sched.JobLocker.Unlock()
    }
}

func (sched *Sched) Fail(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if _, ok := sched.procQueue[jobId]; ok {
        delete(sched.procQueue, jobId)
    }
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) getFuncStat(Func string) *FuncStat {
    defer sched.FuncLocker.Unlock()
    sched.FuncLocker.Lock()
    stat, ok := sched.stats[Func]
    if !ok {
        stat = NewFuncStat(Func)
        sched.stats[Func] = stat
    }
    return stat
}

func (sched *Sched) IncrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Incr()
}

func (sched *Sched) DecrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Decr()
}

func (sched *Sched) IncrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Incr()
}

func (sched *Sched) DecrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Decr()
}

func (sched *Sched) IncrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Incr()
    }
}

func (sched *Sched) DecrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Decr()
    }
}

func (sched *Sched) SchedLater(jobId int64, delay int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if _, ok := sched.procQueue[jobId]; ok {
        delete(sched.procQueue, jobId)
    }
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    var now = time.Now()
    job.SchedAt = int64(now.Unix()) + delay
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) pushJobPQ(job driver.Job) bool {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_READY {
        item := &Item{
            value: job.Id,
            priority: job.SchedAt,
        }
        if sched.cacheItem != nil && item.priority < sched.cacheItem.priority {
            if job.Id == sched.cacheItem.value {
                return true
            }
            job, _ = sched.driver.Get(sched.cacheItem.value)
            sched.cacheItem = item
            if job.Id <= 0 || job.Status != driver.JOB_STATUS_READY {
                return false
            }
        }
        pq, ok := sched.jobPQ[job.Func]
        if !ok {
            pq1 := make(PriorityQueue, 0)
            pq = &pq1
            sched.jobPQ[job.Func] = pq
            heap.Init(pq)
        }
        heap.Push(pq, item)
        return true
    }
    return false
}

func (sched *Sched) pushRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        runAt := job.RunAt
        if runAt == 0 {
            runAt = job.SchedAt
        }
        item := &Item{
            value: job.Id,
            priority: runAt + job.Timeout,
        }
        heap.Push(&sched.revertPQ, item)
    }
}

func (sched *Sched) removeRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        for _, item := range sched.revertPQ {
            if item.value == job.Id {
                heap.Remove(&sched.revertPQ, item.index)
                break
            }
        }
    }
}

func (sched *Sched) loadJobQueue() {
    updateQueue := make([]driver.Job, 0)
    removeQueue := make([]driver.Job, 0)
    var now = time.Now()
    current := int64(now.Unix())

    iter := sched.driver.NewIterator(nil)
    for {
        if !iter.Next() {
            break
        }
        job := iter.Value()
        if job.Name == "" {
            removeQueue = append(removeQueue, job)
            continue
        }
        sched.IncrStatJob(job)
        sched.pushJobPQ(job)
        runAt := job.RunAt
        if runAt < job.SchedAt {
            runAt = job.SchedAt
        }
        if runAt + job.Timeout < current {
            updateQueue = append(updateQueue, job)
        } else {
            sched.JobLocker.Lock()
            sched.procQueue[job.Id] = job
            sched.JobLocker.Unlock()
            sched.IncrStatProc(job)
            sched.pushRevertPQ(job)
        }
    }

    iter.Close()

    for _, job := range updateQueue {
        job.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&job)
    }

    for _, job := range removeQueue {
        sched.driver.Delete(job.Id)
    }
}

func (sched *Sched) Close() {
    sched.alive = false
    sched.driver.Close()
    log.Printf("Periodic task system shutdown\n")
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ezcx

import (
	"context"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	ServerDefaultSignals []os.Signal = []os.Signal{
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGHUP,
	}
)

//Handler interface: Defines the contract.
type Handler interface {
	http.Handler
	Handle(res *WebhookResponse, req *WebhookRequest) error
}

// Functional Adapter: HandlerFunc is an adapter.
// HandlerFunc satisfies the Handler interface
type HandlerFunc func(*WebhookResponse, *WebhookRequest) error

// Seems  redundant; may serve a purpose, though, for structural handlers.
// (ie Need to implement for functional handler to satisfy Handle which would
// require implementation for structural handlers.)
func (h HandlerFunc) Handle(res *WebhookResponse, req *WebhookRequest) error {
	return h(res, req)
}

// yaquino@2022-10-07: http.Request's context is flowd down to the WebhookRequest
// via WebhookRequestFromRequest (requests.go)
func (h HandlerFunc) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()
	req, err := WebhookRequestFromRequest(r)
	if err != nil {
		log.Println("Error during WebhookRequestFromRequest")
		log.Println(err)
		return
	}
	req.ctx = r.Context // flowing down the requests's Context.
	res := req.InitializeResponse()
	err = h(res, req)
	if err != nil {
		log.Println("Error during HandlerFunc execution")
		log.Println(err)
		return
	}
	err = res.WriteResponse(w)
	if err != nil {
		log.Println("Error during WebhookResponse.WriteResponse")
		return
	}
}

type Server struct {
	signals []os.Signal
	signal  chan os.Signal
	errs    chan error
	server  *http.Server
	mux     *http.ServeMux
	lg      *log.Logger
}

func NewServer(ctx context.Context, addr string, lg *log.Logger, signals ...os.Signal) *Server {
	return new(Server).Init(ctx, addr, lg, signals...)
}

func (s *Server) Init(ctx context.Context, addr string, lg *log.Logger, signals ...os.Signal) *Server {
	if len(signals) == 0 {
		s.signals = ServerDefaultSignals
	} else {
		// rethink this later on.  We need to make sure there at least
		// the right group of signals!
		s.signals = signals
	}
	s.signal = make(chan os.Signal, 1)
	signal.Notify(s.signal, s.signals...)

	if lg == nil {
		lg = log.Default()
	}
	s.lg = lg

	s.errs = make(chan error)
	s.mux = http.NewServeMux()
	s.server = &http.Server{
		Addr:        addr,
		Handler:     s.mux,
		BaseContext: func(l net.Listener) context.Context { return ctx },
	}
	return s
}

func (s *Server) SetHandler(h http.Handler) {
	s.server.Handler = h
	if s.isMux(h) {
		s.mux = h.(*http.ServeMux)
	} else {
		s.mux = nil
	}
}

func (s *Server) ServeMux() *http.ServeMux {
	return s.mux
}

func (s *Server) isMux(h http.Handler) bool {
	_, ok := h.(*http.ServeMux)
	return ok
}

func (s *Server) HandleCx(pattern string, handler HandlerFunc) {
	s.mux.Handle(pattern, handler)
}

// yaquino@2022-09-21: I have concerns that checking the parent context will not work as desired.
func (s *Server) ListenAndServe(ctx context.Context) {
	defer func() {
		close(s.errs)
		close(s.signal)
	}()
	// Run ListenAndServe on a separate goroutine.
	s.lg.Printf("EZCX server listening and serving on %s\n", s.server.Addr)
	go func() {
		err := s.server.ListenAndServe()
		if err != nil {
			s.lg.Println(err)
			s.errs <- err
			close(s.errs)
		}
	}()

	for {
		select {
		// If the context is done, we need to return.
		case <-ctx.Done():
			s.lg.Println("EZCX server context is done")
			err := ctx.Err()
			if err != nil {
				s.lg.Print("EZCX server context error...")
				s.lg.Println(err)
			}
			return
		// If there's a non-nil error, we need to return
		case err := <-s.errs:
			if err != nil {
				s.lg.Print("EZCX server non-nil error...")
				s.lg.Println(err)
				return
			}
		case sig := <-s.signal:
			s.lg.Printf("EZCX server signal %s received...", sig)
			switch sig {
			case syscall.SIGHUP:
				s.lg.Println("EZCX reconfigure", sig)
				err := s.Reconfigure()
				if err != nil {
					s.errs <- err
				}
			default:
				s.lg.Printf("EZCX graceful shutdown initiated...")
				err := s.Shutdown(ctx)
				if err != nil {
					s.lg.Println(err)
				}
				s.lg.Println("EZCX shutdown SUCCESS")
				return
			}
		}
	}
}

// Omitted for now.
func (s *Server) Reconfigure() error {
	return nil
}

func (s *Server) Shutdown(ctx context.Context) error {
	timeout, cancel := context.WithTimeout(ctx, time.Second*5)
	defer cancel()
	err := s.server.Shutdown(timeout)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>better error handling<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ezcx

import (
	"context"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	ServerDefaultSignals []os.Signal = []os.Signal{
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGHUP,
	}
)

//Handler interface: Defines the contract.
type Handler interface {
	http.Handler
	Handle(res *WebhookResponse, req *WebhookRequest) error
}

// Functional Adapter: HandlerFunc is an adapter.
// HandlerFunc satisfies the Handler interface
type HandlerFunc func(*WebhookResponse, *WebhookRequest) error

// Seems  redundant; may serve a purpose, though, for structural handlers.
// (ie Need to implement for functional handler to satisfy Handle which would
// require implementation for structural handlers.)
func (h HandlerFunc) Handle(res *WebhookResponse, req *WebhookRequest) error {
	return h(res, req)
}

// yaquino@2022-10-07: http.Request's context is flowd down to the WebhookRequest
// via WebhookRequestFromRequest (requests.go)
func (h HandlerFunc) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()
	req, err := WebhookRequestFromRequest(r)
	if err != nil {
		log.Println("Error during WebhookRequestFromRequest")
		log.Println(err)
		return
	}
	req.ctx = r.Context // flowing down the requests's Context added.
	res := req.InitializeResponse()
	err = h(res, req)
	if err != nil {
		log.Println("Error during HandlerFunc execution")
		log.Println(err)
		return
	}
	err = res.WriteResponse(w)
	if err != nil {
		log.Println("Error during WebhookResponse.WriteResponse")
		return
	}
}

type Server struct {
	signals []os.Signal
	signal  chan os.Signal
	errs    chan error
	server  *http.Server
	mux     *http.ServeMux
	lg      *log.Logger
}

func NewServer(ctx context.Context, addr string, lg *log.Logger, signals ...os.Signal) *Server {
	return new(Server).Init(ctx, addr, lg, signals...)
}

func (s *Server) Init(ctx context.Context, addr string, lg *log.Logger, signals ...os.Signal) *Server {
	if len(signals) == 0 {
		s.signals = ServerDefaultSignals
	} else {
		// rethink this later on.  We need to make sure there at least
		// the right group of signals!
		s.signals = signals
	}
	s.signal = make(chan os.Signal, 1)
	signal.Notify(s.signal, s.signals...)

	if lg == nil {
		lg = log.Default()
	}
	s.lg = lg

	s.errs = make(chan error)
	s.mux = http.NewServeMux()
	s.server = &http.Server{
		Addr:        addr,
		Handler:     s.mux,
		BaseContext: func(l net.Listener) context.Context { return ctx },
	}
	return s
}

func (s *Server) SetHandler(h http.Handler) {
	s.server.Handler = h
	if s.isMux(h) {
		s.mux = h.(*http.ServeMux)
	} else {
		s.mux = nil
	}
}

func (s *Server) ServeMux() *http.ServeMux {
	return s.mux
}

func (s *Server) isMux(h http.Handler) bool {
	_, ok := h.(*http.ServeMux)
	return ok
}

func (s *Server) HandleCx(pattern string, handler HandlerFunc) {
	s.mux.Handle(pattern, handler)
}

// yaquino@2022-09-21: I have concerns that checking the parent context will not work as desired.
func (s *Server) ListenAndServe(ctx context.Context) {
	defer func() {
		close(s.errs)
		close(s.signal)
	}()
	// Run ListenAndServe on a separate goroutine.
	s.lg.Printf("EZCX server listening and serving on %s\n", s.server.Addr)
	go func() {
		err := s.server.ListenAndServe()
		if err != nil {
			s.lg.Println(err)
			s.errs <- err
			close(s.errs)
		}
	}()

	for {
		select {
		// If the context is done, we need to return.
		case <-ctx.Done():
			s.lg.Println("EZCX server context is done")
			err := ctx.Err()
			if err != nil {
				s.lg.Print("EZCX server context error...")
				s.lg.Println(err)
			}
			return
		// If there's a non-nil error, we need to return
		case err := <-s.errs:
			if err != nil {
				s.lg.Print("EZCX server non-nil error...")
				s.lg.Println(err)
				return
			}
		case sig := <-s.signal:
			s.lg.Printf("EZCX server signal %s received...", sig)
			switch sig {
			case syscall.SIGHUP:
				s.lg.Println("EZCX reconfigure", sig)
				err := s.Reconfigure()
				if err != nil {
					s.errs <- err
				}
			default:
				s.lg.Printf("EZCX graceful shutdown initiated...")
				err := s.Shutdown(ctx)
				if err != nil {
					s.lg.Println(err)
				}
				s.lg.Println("EZCX shutdown SUCCESS")
				return
			}
		}
	}
}

// Omitted for now.
func (s *Server) Reconfigure() error {
	return nil
}

func (s *Server) Shutdown(ctx context.Context) error {
	timeout, cancel := context.WithTimeout(ctx, time.Second*5)
	defer cancel()
	err := s.server.Shutdown(timeout)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>