<commit_before>package core

/*
#cgo CFLAGS: -I./src/include
#include "lwip/tcp.h"
*/
import "C"
import (
	"errors"
	"fmt"
	"net"
	"sync"
	"unsafe"
)

type tcpConn struct {
	sync.Mutex

	pcb         *C.struct_tcp_pcb
	handler     ConnectionHandler
	network     string
	remoteAddr  net.Addr
	localAddr   net.Addr
	connKeyArg  unsafe.Pointer
	connKey     uint32
	closing     bool
	localClosed bool
	aborting    bool
	canWrite    *sync.Cond // Condition variable to implement TCP backpressure.
}

func NewTCPConnection(pcb *C.struct_tcp_pcb, handler ConnectionHandler, connKey uint32, connKeyArg unsafe.Pointer) (Connection, error) {
	// Pass the key as arg for subsequent tcp callbacks.
	C.tcp_arg(pcb, unsafe.Pointer(connKeyArg))

	// Register callbacks.
	SetTCPRecvCallback(pcb)
	SetTCPSentCallback(pcb)
	SetTCPErrCallback(pcb)
	SetTCPPollCallback(pcb, C.u8_t(TCP_POLL_INTERVAL))

	conn := &tcpConn{
		pcb:         pcb,
		handler:     handler,
		network:     "tcp",
		localAddr:   ParseTCPAddr(IPAddrNTOA(pcb.remote_ip), uint16(pcb.remote_port)),
		remoteAddr:  ParseTCPAddr(IPAddrNTOA(pcb.local_ip), uint16(pcb.local_port)),
		connKeyArg:  connKeyArg,
		connKey:     connKey,
		closing:     false,
		localClosed: false,
		aborting:    false,
		canWrite:    sync.NewCond(&sync.Mutex{}),
	}

	// Unlocks lwip thread during connecting remote host, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := handler.Connect(conn, conn.RemoteAddr())
	lwipMutex.Lock()
	if err != nil {
		return nil, err
	}
	return conn, nil
}

func (conn *tcpConn) RemoteAddr() net.Addr {
	return conn.remoteAddr
}

func (conn *tcpConn) LocalAddr() net.Addr {
	return conn.localAddr
}

func (conn *tcpConn) Receive(data []byte) error {
	if conn.isClosing() {
		return errors.New(fmt.Sprintf("connection %v->%v was closed by remote", conn.LocalAddr(), conn.RemoteAddr()))
	}
	if conn.isAborting() {
		return errors.New(fmt.Sprintf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr()))
	}
	// Unlocks lwip thread during sending data to remote, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := conn.handler.DidReceive(conn, data)
	lwipMutex.Lock()
	if err != nil {
		return errors.New(fmt.Sprintf("write proxy failed: %v", err))
	}
	C.tcp_recved(conn.pcb, C.u16_t(len(data)))
	return nil
}

// tcpWrite enqueues data to snd_buf, and treats ERR_MEM returned by tcp_write not an error,
// but instead tells the caller that data is not successfully enqueued, and should try
// again another time. By calling this function, the lwIP thread is assumed to be already
// locked by the caller.
func (conn *tcpConn) tcpWrite(data []byte) (int, error) {
	err := C.tcp_write(conn.pcb, unsafe.Pointer(&data[0]), C.u16_t(len(data)), C.TCP_WRITE_FLAG_COPY)
	if err == C.ERR_OK {
		C.tcp_output(conn.pcb)
		return len(data), nil
	} else if err == C.ERR_MEM {
		return 0, nil
	}
	return 0, fmt.Errorf("lwip tcp_write failed with error code: %v", int(err))
}

func (conn *tcpConn) Write(data []byte) (int, error) {
	totalWritten := 0

	conn.canWrite.L.Lock()
	defer conn.canWrite.L.Unlock()

	for len(data) > 0 {
		if conn.isLocalClosed() {
			return totalWritten, fmt.Errorf("connection %v->%v was closed by local", conn.LocalAddr(), conn.RemoteAddr())
		}
		if conn.isAborting() {
			return totalWritten, fmt.Errorf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr())
		}

		lwipMutex.Lock()
		toWrite := len(data)
		if toWrite > int(conn.pcb.snd_buf) {
			// Write at most the size of the LWIP buffer.
			toWrite = int(conn.pcb.snd_buf)
		}
		if toWrite > 0 {
			written, err := conn.tcpWrite(data[0:toWrite])
			totalWritten += written
			if err != nil {
				lwipMutex.Unlock()
				return totalWritten, err
			}
			data = data[written:len(data)]
		}
		lwipMutex.Unlock()
		if len(data) == 0 {
			break // Don't block if all the data has been written.
		}
		conn.canWrite.Wait()
	}

	return totalWritten, nil
}

func (conn *tcpConn) Sent(len uint16) error {
	conn.handler.DidSend(conn, len)
	// Some packets are acknowledged by local client, check if any pending data to send.
	return conn.CheckState()
}

func (conn *tcpConn) isClosing() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.closing
}

func (conn *tcpConn) isAborting() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.aborting
}

func (conn *tcpConn) isLocalClosed() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.localClosed
}

func (conn *tcpConn) CheckState() error {
	if conn.isAborting() {
		conn.abortInternal()
		return NewLWIPError(LWIP_ERR_ABRT)
	}

	if conn.isClosing() || conn.isLocalClosed() {
		conn.closeInternal()
		return NewLWIPError(LWIP_ERR_OK)
	}

	// Signal the writer to try writting.
	conn.canWrite.Broadcast()

	return NewLWIPError(LWIP_ERR_OK)
}

func (conn *tcpConn) Close() error {
	conn.Lock()
	defer conn.Unlock()

	// Close maybe called outside of lwIP thread, we should not call tcp_close() in this
	// function, instead just make a flag to indicate we are closing the connection.
	conn.closing = true
	return nil
}

func (conn *tcpConn) setLocalClosed() error {
	conn.Lock()
	defer conn.Unlock()

	conn.localClosed = true
	conn.canWrite.Broadcast()
	return nil
}

func (conn *tcpConn) closeInternal() error {
	C.tcp_arg(conn.pcb, nil)
	C.tcp_recv(conn.pcb, nil)
	C.tcp_sent(conn.pcb, nil)
	C.tcp_err(conn.pcb, nil)
	C.tcp_poll(conn.pcb, nil, 0)

	conn.Release()

	// TODO: may return ERR_MEM if no memory to allocate segments use for closing the conn,
	// should check and try again in Sent() for Poll() callbacks.
	err := C.tcp_close(conn.pcb)
	if err == C.ERR_OK {
		return nil
	} else {
		return errors.New(fmt.Sprintf("close TCP connection failed, lwip error code %d", int(err)))
	}
}

func (conn *tcpConn) abortInternal() {
	conn.Release()
	C.tcp_abort(conn.pcb)
}

func (conn *tcpConn) Abort() {
	conn.Lock()
	defer conn.Unlock()

	conn.aborting = true
	conn.canWrite.Broadcast()
}

// The corresponding pcb is already freed when this callback is called
func (conn *tcpConn) Err(err error) {
	conn.Release()
	conn.handler.DidClose(conn)
}

func (conn *tcpConn) LocalDidClose() error {
	conn.handler.LocalDidClose(conn)
	conn.setLocalClosed()    // flag closing
	return conn.CheckState() // check pending data
}

func (conn *tcpConn) Release() {
	if _, found := tcpConns.Load(conn.connKey); found {
		FreeConnKeyArg(conn.connKeyArg)
		tcpConns.Delete(conn.connKey)
	}
}

func (conn *tcpConn) Poll() error {
	return conn.CheckState()
}
<commit_msg>close the tx side as soon as possible<commit_after>package core

/*
#cgo CFLAGS: -I./src/include
#include "lwip/tcp.h"
*/
import "C"
import (
	"errors"
	"fmt"
	"net"
	"sync"
	"unsafe"
)

type tcpConn struct {
	sync.Mutex

	pcb         *C.struct_tcp_pcb
	handler     ConnectionHandler
	network     string
	remoteAddr  net.Addr
	localAddr   net.Addr
	connKeyArg  unsafe.Pointer
	connKey     uint32
	closing     bool
	localClosed bool
	aborting    bool
	canWrite    *sync.Cond // Condition variable to implement TCP backpressure.
}

func NewTCPConnection(pcb *C.struct_tcp_pcb, handler ConnectionHandler, connKey uint32, connKeyArg unsafe.Pointer) (Connection, error) {
	// Pass the key as arg for subsequent tcp callbacks.
	C.tcp_arg(pcb, unsafe.Pointer(connKeyArg))

	// Register callbacks.
	SetTCPRecvCallback(pcb)
	SetTCPSentCallback(pcb)
	SetTCPErrCallback(pcb)
	SetTCPPollCallback(pcb, C.u8_t(TCP_POLL_INTERVAL))

	conn := &tcpConn{
		pcb:         pcb,
		handler:     handler,
		network:     "tcp",
		localAddr:   ParseTCPAddr(IPAddrNTOA(pcb.remote_ip), uint16(pcb.remote_port)),
		remoteAddr:  ParseTCPAddr(IPAddrNTOA(pcb.local_ip), uint16(pcb.local_port)),
		connKeyArg:  connKeyArg,
		connKey:     connKey,
		closing:     false,
		localClosed: false,
		aborting:    false,
		canWrite:    sync.NewCond(&sync.Mutex{}),
	}

	// Unlocks lwip thread during connecting remote host, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := handler.Connect(conn, conn.RemoteAddr())
	lwipMutex.Lock()
	if err != nil {
		return nil, err
	}
	return conn, nil
}

func (conn *tcpConn) RemoteAddr() net.Addr {
	return conn.remoteAddr
}

func (conn *tcpConn) LocalAddr() net.Addr {
	return conn.localAddr
}

func (conn *tcpConn) Receive(data []byte) error {
	if conn.isClosing() {
		return errors.New(fmt.Sprintf("connection %v->%v was closed by remote", conn.LocalAddr(), conn.RemoteAddr()))
	}
	if conn.isAborting() {
		return errors.New(fmt.Sprintf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr()))
	}
	// Unlocks lwip thread during sending data to remote, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := conn.handler.DidReceive(conn, data)
	lwipMutex.Lock()
	if err != nil {
		return errors.New(fmt.Sprintf("write proxy failed: %v", err))
	}
	C.tcp_recved(conn.pcb, C.u16_t(len(data)))
	return nil
}

// tcpWrite enqueues data to snd_buf, and treats ERR_MEM returned by tcp_write not an error,
// but instead tells the caller that data is not successfully enqueued, and should try
// again another time. By calling this function, the lwIP thread is assumed to be already
// locked by the caller.
func (conn *tcpConn) tcpWrite(data []byte) (int, error) {
	err := C.tcp_write(conn.pcb, unsafe.Pointer(&data[0]), C.u16_t(len(data)), C.TCP_WRITE_FLAG_COPY)
	if err == C.ERR_OK {
		C.tcp_output(conn.pcb)
		return len(data), nil
	} else if err == C.ERR_MEM {
		return 0, nil
	}
	return 0, fmt.Errorf("lwip tcp_write failed with error code: %v", int(err))
}

func (conn *tcpConn) Write(data []byte) (int, error) {
	totalWritten := 0

	conn.canWrite.L.Lock()
	defer conn.canWrite.L.Unlock()

	for len(data) > 0 {
		if conn.isLocalClosed() {
			return totalWritten, fmt.Errorf("connection %v->%v was closed by local", conn.LocalAddr(), conn.RemoteAddr())
		}
		if conn.isAborting() {
			return totalWritten, fmt.Errorf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr())
		}

		lwipMutex.Lock()
		toWrite := len(data)
		if toWrite > int(conn.pcb.snd_buf) {
			// Write at most the size of the LWIP buffer.
			toWrite = int(conn.pcb.snd_buf)
		}
		if toWrite > 0 {
			written, err := conn.tcpWrite(data[0:toWrite])
			totalWritten += written
			if err != nil {
				lwipMutex.Unlock()
				return totalWritten, err
			}
			data = data[written:len(data)]
		}
		lwipMutex.Unlock()
		if len(data) == 0 {
			break // Don't block if all the data has been written.
		}
		conn.canWrite.Wait()
	}

	return totalWritten, nil
}

func (conn *tcpConn) Sent(len uint16) error {
	conn.handler.DidSend(conn, len)
	// Some packets are acknowledged by local client, check if any pending data to send.
	return conn.CheckState()
}

func (conn *tcpConn) isClosing() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.closing
}

func (conn *tcpConn) isAborting() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.aborting
}

func (conn *tcpConn) isLocalClosed() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.localClosed
}

func (conn *tcpConn) CheckState() error {
	if conn.isAborting() {
		conn.abortInternal()
		return NewLWIPError(LWIP_ERR_ABRT)
	}

	if conn.isClosing() || conn.isLocalClosed() {
		conn.closeInternal()
		return NewLWIPError(LWIP_ERR_OK)
	}

	// Signal the writer to try writting.
	conn.canWrite.Broadcast()

	return NewLWIPError(LWIP_ERR_OK)
}

func (conn *tcpConn) Close() error {
	lwipMutex.Lock()
	C.tcp_shutdown(conn.pcb, 0, 1) // Close the TX side ASAP.
	lwipMutex.Unlock()

	conn.Lock()
	defer conn.Unlock()

	// Close maybe called outside of lwIP thread, we should not call tcp_close() in this
	// function, instead just make a flag to indicate we are closing the connection.
	conn.closing = true
	return nil
}

func (conn *tcpConn) setLocalClosed() error {
	conn.Lock()
	defer conn.Unlock()

	conn.localClosed = true
	conn.canWrite.Broadcast()
	return nil
}

func (conn *tcpConn) closeInternal() error {
	C.tcp_arg(conn.pcb, nil)
	C.tcp_recv(conn.pcb, nil)
	C.tcp_sent(conn.pcb, nil)
	C.tcp_err(conn.pcb, nil)
	C.tcp_poll(conn.pcb, nil, 0)

	conn.Release()

	// TODO: may return ERR_MEM if no memory to allocate segments use for closing the conn,
	// should check and try again in Sent() for Poll() callbacks.
	err := C.tcp_close(conn.pcb)
	if err == C.ERR_OK {
		return nil
	} else {
		return errors.New(fmt.Sprintf("close TCP connection failed, lwip error code %d", int(err)))
	}
}

func (conn *tcpConn) abortInternal() {
	conn.Release()
	C.tcp_abort(conn.pcb)
}

func (conn *tcpConn) Abort() {
	conn.Lock()
	defer conn.Unlock()

	conn.aborting = true
	conn.canWrite.Broadcast()
}

// The corresponding pcb is already freed when this callback is called
func (conn *tcpConn) Err(err error) {
	conn.Release()
	conn.handler.DidClose(conn)
}

func (conn *tcpConn) LocalDidClose() error {
	conn.handler.LocalDidClose(conn)
	conn.setLocalClosed()    // flag closing
	return conn.CheckState() // check pending data
}

func (conn *tcpConn) Release() {
	if _, found := tcpConns.Load(conn.connKey); found {
		FreeConnKeyArg(conn.connKeyArg)
		tcpConns.Delete(conn.connKey)
	}
}

func (conn *tcpConn) Poll() error {
	return conn.CheckState()
}
<|endoftext|>


<commit_before>package core

import (
	"bufio"
	"log"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/xlvector/hector/util"
)

type CombinedFeature []string

type FeatureSplit []float64

func FindCategory(split []float64, value float64) int {
	return sort.Search(len(split), func(i int) bool { return split[i] >= value })
}

/* RawDataSet */
type RawDataSet struct {
	Samples     []*RawSample
	FeatureKeys map[string]bool
}

func NewRawDataSet() *RawDataSet {
	ret := RawDataSet{}
	ret.Samples = []*RawSample{}
	ret.FeatureKeys = make(map[string]bool)
	return &ret
}

func (d *RawDataSet) AddSample(sample *RawSample) {
	d.Samples = append(d.Samples, sample)
}

func (d *RawDataSet) ToDataSet(splits map[string][]float64, combinations []CombinedFeature) *DataSet {
	out_data := NewDataSet()
	fm := make(map[string]int64)
	for _, sample := range d.Samples {
		out_sample := NewSample()
		out_sample.Label = sample.Label
		if splits != nil {
			for fkey_str, fvalue_str := range sample.Features {
				fkey := ""
				fvalue := 0.0
				if GetFeatureType(fkey_str) == FeatureTypeEnum.CONTINUOUS_FEATURE {
					split, ok := splits[fkey_str]
					if ok {
						cat := FindCategory(split, util.ParseFloat64(fvalue_str))
						fkey = fkey_str + "_" + strconv.FormatInt(int64(cat), 10)
						fvalue = 1.0
					} else {
						fvalue = util.ParseFloat64(fvalue_str)
					}
					fm[fkey] = util.Hash(fkey)
					out_sample.AddFeature(Feature{Id: util.Hash(fkey), Value: fvalue})
				}
			}
		}
		for _, combination := range combinations {
			fkey := ""
			for _, ckey := range combination {
				fkey += ckey
				fkey += ":"
				fkey += sample.GetFeatureValue(ckey)
				fkey += "_"
			}
			fm[fkey] = util.Hash(fkey)
			out_sample.AddFeature(Feature{Id: util.Hash(fkey), Value: 1.0})
		}
		out_data.AddSample(out_sample)
	}
	f, _ := os.Create("features.tsv")
	defer f.Close()
	w := bufio.NewWriter(f)
	for k, v := range fm {
		w.WriteString(k + "\t" + strconv.FormatInt(v, 10) + "\n")
	}

	return out_data
}

func (d *RawDataSet) Load(path string) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	ch := make(chan string, 1000)
	go func() {
		reader := bufio.NewReader(file)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				break
			}
			ch <- line
		}
		close(ch)
	}()

	for line := range ch {
		line = strings.Replace(line, " ", "\t", -1)
		tks := strings.Split(line, "\t")
		sample := NewRawSample()
		for i, tk := range tks {
			if i == 0 {
				label, err := strconv.ParseInt(tk, 10, 16)
				if err != nil {
					break
				}
				if label > 0 {
					sample.Label = 1.0
				} else {
					sample.Label = 0.0
				}
			} else {
				kv := strings.Split(tk, ":")
				sample.Features[kv[0]] = kv[1]
				d.FeatureKeys[kv[0]] = true
			}
		}
		d.AddSample(sample)
	}
	return nil
}

/* DataSet */
type DataSet struct {
	Samples          []*Sample
	FeatureNameIdMap map[int64]string
	max_label        int
}

func NewDataSet() *DataSet {
	ret := DataSet{}
	ret.Samples = []*Sample{}
	ret.FeatureNameIdMap = make(map[int64]string)
	return &ret
}

func (d *DataSet) AddSample(sample *Sample) {
	d.Samples = append(d.Samples, sample)
	if d.max_label < sample.Label {
		d.max_label = sample.Label
	}
}

func (d *DataSet) Load(path string, global_bias_feature_id int64) error {
	fm := make(map[string]int64)
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := strings.Replace(scanner.Text(), " ", "\t", -1)
		tks := strings.Split(line, "\t")
		sample := Sample{Features: []Feature{}, Label: 0}
		for i, tk := range tks {
			if i == 0 {
				label, _ := strconv.Atoi(tk)
				sample.Label = label
				if d.max_label < label {
					d.max_label = label
				}
			} else {
				kv := strings.Split(tk, ":")
				feature_id, err := strconv.ParseInt(kv[0], 10, 64)
				if err != nil {
					feature_id = util.Hash(kv[0])
					fm[kv[0]] = feature_id
				}
				d.FeatureNameIdMap[feature_id] = kv[0]
				feature_value := 1.0
				if len(kv) > 1 {
					feature_value, err = strconv.ParseFloat(kv[1], 64)
					if err != nil {
						break
					}
				}
				feature := Feature{feature_id, feature_value}
				sample.Features = append(sample.Features, feature)
			}
		}
		if global_bias_feature_id >= 0 {
			sample.Features = append(sample.Features, Feature{global_bias_feature_id, 1.0})
		}
		d.AddSample(&sample)
	}
	if scanner.Err() != nil {
		return scanner.Err()
	}
	f, _ := os.Create("features.tsv")
	defer f.Close()
	w := bufio.NewWriter(f)
	for k, v := range fm {
		w.WriteString(k + "\t" + strconv.FormatInt(v, 10) + "\n")
	}

	log.Println("dataset size : ", len(d.Samples))
	return nil
}

func RemoveLowFreqFeatures(dataset *DataSet, threshold float64) {
	freq := NewVector()

	for _, sample := range dataset.Samples {
		for _, feature := range sample.Features {
			freq.AddValue(feature.Id, 1.0)
		}
	}

	for _, sample := range dataset.Samples {
		features := []Feature{}
		for _, feature := range sample.Features {
			if freq.GetValue(feature.Id) > threshold {
				features = append(features, feature)
			}
		}
		sample.Features = features
	}
}

func (d *DataSet) Split(f func(int) bool) *DataSet {
	out_data := NewDataSet()
	for i, sample := range d.Samples {
		if f(i) {
			out_data.AddSample(sample)
		}
	}
	return out_data
}

/* Real valued DataSet */
type RealDataSet struct {
	Samples []*RealSample
}

func NewRealDataSet() *RealDataSet {
	ret := RealDataSet{}
	ret.Samples = []*RealSample{}
	return &ret
}

func (d *RealDataSet) AddSample(sample *RealSample) {
	d.Samples = append(d.Samples, sample)
}

func (d *RealDataSet) Load(path string, global_bias_feature_id int64) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := strings.Replace(scanner.Text(), " ", "\t", -1)
		tks := strings.Split(line, "\t")
		sample := RealSample{Features: []Feature{}, Value: 0.0}
		for i, tk := range tks {
			if i == 0 {
				value := util.ParseFloat64(tk)
				sample.Value = value
			} else {
				kv := strings.Split(tk, ":")
				feature_id, err := strconv.ParseInt(kv[0], 10, 64)
				if err != nil {
					break
				}
				feature_value := 1.0
				if len(kv) > 1 {
					feature_value, err = strconv.ParseFloat(kv[1], 64)
					if err != nil {
						break
					}
				}
				feature := Feature{feature_id, feature_value}
				sample.Features = append(sample.Features, feature)
			}
		}
		if global_bias_feature_id >= 0 {
			sample.Features = append(sample.Features, Feature{global_bias_feature_id, 1.0})
		}
		d.AddSample(&sample)
	}
	if scanner.Err() != nil {
		return scanner.Err()
	}
	return nil
}
<commit_msg>read file in go routine<commit_after>package core

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/xlvector/hector/util"
)

type CombinedFeature []string

type FeatureSplit []float64

func FindCategory(split []float64, value float64) int {
	return sort.Search(len(split), func(i int) bool { return split[i] >= value })
}

/* RawDataSet */
type RawDataSet struct {
	Samples     []*RawSample
	FeatureKeys map[string]bool
}

func NewRawDataSet() *RawDataSet {
	ret := RawDataSet{}
	ret.Samples = []*RawSample{}
	ret.FeatureKeys = make(map[string]bool)
	return &ret
}

func (d *RawDataSet) AddSample(sample *RawSample) {
	d.Samples = append(d.Samples, sample)
}

func (d *RawDataSet) ToDataSet(splits map[string][]float64, combinations []CombinedFeature) *DataSet {
	out_data := NewDataSet()
	fm := make(map[string]int64)
	for _, sample := range d.Samples {
		out_sample := NewSample()
		out_sample.Label = sample.Label
		if splits != nil {
			for fkey_str, fvalue_str := range sample.Features {
				fkey := ""
				fvalue := 0.0
				if GetFeatureType(fkey_str) == FeatureTypeEnum.CONTINUOUS_FEATURE {
					split, ok := splits[fkey_str]
					if ok {
						cat := FindCategory(split, util.ParseFloat64(fvalue_str))
						fkey = fkey_str + "_" + strconv.FormatInt(int64(cat), 10)
						fvalue = 1.0
					} else {
						fvalue = util.ParseFloat64(fvalue_str)
					}
					fm[fkey] = util.Hash(fkey)
					out_sample.AddFeature(Feature{Id: util.Hash(fkey), Value: fvalue})
				}
			}
		}
		for _, combination := range combinations {
			fkey := ""
			for _, ckey := range combination {
				fkey += ckey
				fkey += ":"
				fkey += sample.GetFeatureValue(ckey)
				fkey += "_"
			}
			fm[fkey] = util.Hash(fkey)
			out_sample.AddFeature(Feature{Id: util.Hash(fkey), Value: 1.0})
		}
		out_data.AddSample(out_sample)
	}
	f, _ := os.Create("features.tsv")
	defer f.Close()
	w := bufio.NewWriter(f)
	for k, v := range fm {
		w.WriteString(k + "\t" + strconv.FormatInt(v, 10) + "\n")
	}

	return out_data
}

func (d *RawDataSet) Load(path string) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	ch := make(chan string, 1000)
	go func() {
		reader := bufio.NewReader(file)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				break
			}
			ch <- line
		}
		close(ch)
	}()

	n := 0
	for line := range ch {
		n += 1
		if n%10000 == 0 {
			fmt.Println(n, len(ch))
		}
		line = strings.Replace(line, " ", "\t", -1)
		tks := strings.Split(line, "\t")
		sample := NewRawSample()
		for i, tk := range tks {
			if i == 0 {
				label, err := strconv.ParseInt(tk, 10, 16)
				if err != nil {
					break
				}
				if label > 0 {
					sample.Label = 1.0
				} else {
					sample.Label = 0.0
				}
			} else {
				kv := strings.Split(tk, ":")
				sample.Features[kv[0]] = kv[1]
				d.FeatureKeys[kv[0]] = true
			}
		}
		d.AddSample(sample)
	}
	return nil
}

/* DataSet */
type DataSet struct {
	Samples          []*Sample
	FeatureNameIdMap map[int64]string
	max_label        int
}

func NewDataSet() *DataSet {
	ret := DataSet{}
	ret.Samples = []*Sample{}
	ret.FeatureNameIdMap = make(map[int64]string)
	return &ret
}

func (d *DataSet) AddSample(sample *Sample) {
	d.Samples = append(d.Samples, sample)
	if d.max_label < sample.Label {
		d.max_label = sample.Label
	}
}

func (d *DataSet) Load(path string, global_bias_feature_id int64) error {
	fm := make(map[string]int64)
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := strings.Replace(scanner.Text(), " ", "\t", -1)
		tks := strings.Split(line, "\t")
		sample := Sample{Features: []Feature{}, Label: 0}
		for i, tk := range tks {
			if i == 0 {
				label, _ := strconv.Atoi(tk)
				sample.Label = label
				if d.max_label < label {
					d.max_label = label
				}
			} else {
				kv := strings.Split(tk, ":")
				feature_id, err := strconv.ParseInt(kv[0], 10, 64)
				if err != nil {
					feature_id = util.Hash(kv[0])
					fm[kv[0]] = feature_id
				}
				d.FeatureNameIdMap[feature_id] = kv[0]
				feature_value := 1.0
				if len(kv) > 1 {
					feature_value, err = strconv.ParseFloat(kv[1], 64)
					if err != nil {
						break
					}
				}
				feature := Feature{feature_id, feature_value}
				sample.Features = append(sample.Features, feature)
			}
		}
		if global_bias_feature_id >= 0 {
			sample.Features = append(sample.Features, Feature{global_bias_feature_id, 1.0})
		}
		d.AddSample(&sample)
	}
	if scanner.Err() != nil {
		return scanner.Err()
	}
	f, _ := os.Create("features.tsv")
	defer f.Close()
	w := bufio.NewWriter(f)
	for k, v := range fm {
		w.WriteString(k + "\t" + strconv.FormatInt(v, 10) + "\n")
	}

	log.Println("dataset size : ", len(d.Samples))
	return nil
}

func RemoveLowFreqFeatures(dataset *DataSet, threshold float64) {
	freq := NewVector()

	for _, sample := range dataset.Samples {
		for _, feature := range sample.Features {
			freq.AddValue(feature.Id, 1.0)
		}
	}

	for _, sample := range dataset.Samples {
		features := []Feature{}
		for _, feature := range sample.Features {
			if freq.GetValue(feature.Id) > threshold {
				features = append(features, feature)
			}
		}
		sample.Features = features
	}
}

func (d *DataSet) Split(f func(int) bool) *DataSet {
	out_data := NewDataSet()
	for i, sample := range d.Samples {
		if f(i) {
			out_data.AddSample(sample)
		}
	}
	return out_data
}

/* Real valued DataSet */
type RealDataSet struct {
	Samples []*RealSample
}

func NewRealDataSet() *RealDataSet {
	ret := RealDataSet{}
	ret.Samples = []*RealSample{}
	return &ret
}

func (d *RealDataSet) AddSample(sample *RealSample) {
	d.Samples = append(d.Samples, sample)
}

func (d *RealDataSet) Load(path string, global_bias_feature_id int64) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := strings.Replace(scanner.Text(), " ", "\t", -1)
		tks := strings.Split(line, "\t")
		sample := RealSample{Features: []Feature{}, Value: 0.0}
		for i, tk := range tks {
			if i == 0 {
				value := util.ParseFloat64(tk)
				sample.Value = value
			} else {
				kv := strings.Split(tk, ":")
				feature_id, err := strconv.ParseInt(kv[0], 10, 64)
				if err != nil {
					break
				}
				feature_value := 1.0
				if len(kv) > 1 {
					feature_value, err = strconv.ParseFloat(kv[1], 64)
					if err != nil {
						break
					}
				}
				feature := Feature{feature_id, feature_value}
				sample.Features = append(sample.Features, feature)
			}
		}
		if global_bias_feature_id >= 0 {
			sample.Features = append(sample.Features, Feature{global_bias_feature_id, 1.0})
		}
		d.AddSample(&sample)
	}
	if scanner.Err() != nil {
		return scanner.Err()
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import "errors"

type broadcastViewer struct {
	// This function may return `false` to signal that it cannot write any more data.
	// The stream will resynchronize at next keyframe.
	write func(data []byte) bool
	// Viewers may hop between streams, but should only receive headers once.
	// This includes track info, as codecs must stay the same between segments.
	skipHeaders bool
	// We group blocks into indeterminate-length clusters. So long as
	// the cluster's timecode has not changed, there's no need to start a new one.
	skipCluster bool
	// To avoid decoding errors due to missing reference frames, the first
	// frame of each track received by a viewer must be a keyframe.
	// Each track for which a keyframe has been sent is marked by a bit here.
	seenKeyframes uint32
}

type Broadcast struct {
	Closed   bool // When set to `true`, all viewers receive an empty bytearray as a notification.
	HasVideo bool
	HasAudio bool
	Width    uint // Dimensions of the video track that came last in the `Tracks` tag.
	Height   uint // Hopefully, there's only one video track in the file.

	viewers map[chan<- []byte]*broadcastViewer
	buffer  []byte
	header  []byte // The EBML (DocType) tag.
	tracks  []byte // The beginning of the Segment (Tracks + Info).

	time struct {
		last  uint64 // Last seen block timecode. The next timecode must be no less than that.
		recv  uint64 // Last received cluster timecode, shifted to ensure monotonicity.
		sent  uint64 // Last sent cluster timecode. (All viewers receive same clusters.)
		shift uint64 // By how much the cluster timecode has been shifted.
	}
}

func NewBroadcast() Broadcast {
	return Broadcast{viewers: make(map[chan<- []byte]*broadcastViewer)}
}

func (cast *Broadcast) Close() error {
	cast.Closed = true
	for _, cb := range cast.viewers {
		cb.write([]byte{})
	}
	return nil
}

func (cast *Broadcast) Connect(ch chan<- []byte, skipHeaders bool) {
	write := func(data []byte) bool {
		// `Broadcast.Write` emits data in block-sized chunks.
		// Thus the buffer size is measured in frames, not bytes.
		if len(ch) == cap(ch) {
			return false
		}
		ch <- data
		return true
	}

	cast.viewers[ch] = &broadcastViewer{write, skipHeaders, false, 0}
}

func (cast *Broadcast) Disconnect(ch chan<- []byte) {
	delete(cast.viewers, ch)
}

func (cast *Broadcast) Reset() {
	cast.buffer = nil
}

func (cast *Broadcast) Write(data []byte) (int, error) {
	cast.buffer = append(cast.buffer, data...)

	for {
		buf := cast.buffer
		tag := EBMLParseTagIncomplete(buf)
		if tag.Consumed == 0 {
			return len(data), nil
		}

		if tag.ID == EBMLSegmentTag || tag.ID == EBMLTracksTag || tag.ID == EBMLClusterTag {
			// Parse the contents of these tags in the same loop.
			buf = buf[:tag.Consumed]
			// Chrome crashes if an indeterminate length is not encoded as 0xFF.
			// If we want to recode it, we'll also need some space for a Void tag.
			if tag.Length == EBMLIndeterminate && tag.Consumed >= 7 {
				cast.buffer[4] = 0xFF
				cast.buffer[5] = EBMLVoidTag
				cast.buffer[6] = 0x80 | byte(tag.Consumed-7)
			}
		} else {
			total := tag.Length + uint64(tag.Consumed)
			if total > 1024*1024 {
				return 0, errors.New("data block too big")
			}

			if total > uint64(len(buf)) {
				return len(data), nil
			}

			buf = buf[:total]
		}

		switch tag.ID {
		case EBMLSeekHeadTag:
			// Disallow seeking.
		case EBMLChaptersTag:
			// Disallow seeking again.
		case EBMLCuesTag:
			// Disallow even more seeking.
		case EBMLVoidTag:
			// Waste of space.
		case EBMLTagsTag:
			// Maybe later.
		case EBMLClusterTag:
			// Ignore boundaries, we'll regroup the data anyway.
		case EBMLPrevSizeTag:
			// Disallow backward seeking too.

		case EBMLEBMLTag:
			// The header is the same in all WebM-s.
			if len(cast.header) == 0 {
				cast.header = append([]byte{}, buf...)
			}

		case EBMLSegmentTag:
			cast.HasVideo = false
			cast.HasAudio = false
			cast.Width = 0
			cast.Height = 0
			cast.tracks = append([]byte{}, buf...)
			// Will recalculate this when the first block arrives.
			cast.time.shift = 0

		case EBMLInfoTag:
			// Default timecode resolution in Matroska is 1 ms. This value is required
			// in WebM; we'll check just in case. Obviously, our timecode rewriting
			// logic won't work with non-millisecond resolutions.
			var scale uint64 = 0

			for buf2 := tag.Contents(buf); len(buf2) != 0; {
				tag2 := EBMLParseTag(buf2)

				switch tag2.ID {
				case 0:
					return 0, errors.New("malformed EBML")

				case EBMLDurationTag:
					total := tag2.Length + uint64(tag2.Consumed) - 2
					if total > 0x7F {
						// I'd rather avoid shifting memory. What kind of integer
						// needs 128 bytes, anyway?
						return 0, errors.New("EBML Duration too large")
					}
					// Live streams must not have a duration.
					buf2[0] = EBMLVoidTag
					buf2[1] = 0x80 | byte(total)

				case EBMLTimecodeScaleTag:
					scale = EBMLParseFixedUint(tag2.Contents(buf2))
				}

				buf2 = tag2.Skip(buf2)
			}

			if scale != 1000000 {
				return 0, errors.New("invalid timecode scale")
			}

			cast.tracks = append(cast.tracks, buf...)

		case EBMLTrackEntryTag:
			// Since `viewer.seenKeyframes` is a 32-bit vector,
			// we need to check that there are at most 32 tracks.
			for buf2 := tag.Contents(buf); len(buf2) != 0; {
				tag2 := EBMLParseTag(buf2)

				switch tag2.ID {
				case 0:
					return 0, errors.New("malformed EBML")

				case EBMLTrackNumberTag:
					// go needs sizeof.
					if t := EBMLParseFixedUint(tag2.Contents(buf2)); t >= 32 {
						return 0, errors.New("too many tracks?")
					}

				case EBMLAudioTag:
					cast.HasAudio = true

				case EBMLVideoTag:
					cast.HasVideo = true
					// While we're here, let's grab some metadata, too.
					for buf3 := tag2.Contents(buf2); len(buf3) != 0; {
						tag3 := EBMLParseTag(buf3)

						switch tag3.ID {
						case 0:
							return 0, errors.New("malformed EBML")

						case EBMLPixelWidthTag:
							cast.Width = uint(EBMLParseFixedUint(tag3.Contents(buf3)))

						case EBMLPixelHeightTag:
							cast.Height = uint(EBMLParseFixedUint(tag3.Contents(buf3)))
						}

						buf3 = tag3.Skip(buf3)
					}
				}

				buf2 = tag2.Skip(buf2)
			}

			cast.tracks = append(cast.tracks, buf...)

		case EBMLTracksTag:
			cast.tracks = append(cast.tracks, buf...)

		case EBMLTimecodeTag:
			// Will reencode it when sending a Cluster.
			cast.time.recv = EBMLParseFixedUint(tag.Contents(buf)) + cast.time.shift

		case EBMLBlockGroupTag, EBMLSimpleBlockTag:
			key := false
			block := tag.Contents(buf)

			if tag.ID == EBMLBlockGroupTag {
				key, block = true, nil

				for buf2 := tag.Contents(buf); len(buf2) != 0; {
					tag2 := EBMLParseTag(buf2)

					switch tag2.ID {
					case 0:
						return 0, errors.New("malformed EBML")

					case EBMLBlockTag:
						block = tag2.Contents(buf2)

					case EBMLReferenceBlockTag:
						// Keyframes, by definition, have no reference frame.
						key = EBMLParseFixedUint(tag2.Contents(buf2)) == 0
					}

					buf2 = tag2.Skip(buf2)
				}

				if block == nil {
					return 0, errors.New("a BlockGroup contains no Blocks")
				}
			}

			track := EBMLParseUint(block)
			if track.Consumed == 0 || track.Value >= 32 || len(block) < track.Consumed+3 {
				return 0, errors.New("invalid track")
			}
			// This bit is always 0 in a Block, but 1 in a keyframe SimpleBlock.
			key = key || block[track.Consumed+2]&0x80 != 0
			// Block timecodes are relative to cluster ones.
			timecode := uint64(block[track.Consumed+0])<<8 | uint64(block[track.Consumed+1])
			if cast.time.recv+timecode < cast.time.last {
				cast.time.shift += cast.time.last - (cast.time.recv + timecode)
				cast.time.recv = cast.time.last - timecode
			}
			cast.time.last = cast.time.recv + timecode

			ctc := cast.time.recv
			cluster := []byte{
				EBMLClusterTag >> 24 & 0xFF,
				EBMLClusterTag >> 16 & 0xFF,
				EBMLClusterTag >> 8 & 0xFF,
				EBMLClusterTag & 0xFF, 0xFF,
				EBMLTimecodeTag, 0x88,
				byte(ctc >> 56), byte(ctc >> 48), byte(ctc >> 40), byte(ctc >> 32),
				byte(ctc >> 24), byte(ctc >> 16), byte(ctc >> 8), byte(ctc),
			}

			trackMask := uint32(1) << track.Value
			for _, cb := range cast.viewers {
				if !cb.skipHeaders {
					if !cb.write(cast.header) || !cb.write(cast.tracks) {
						continue
					}

					cb.skipHeaders = true
					cb.skipCluster = false
				}

				if key {
					cb.seenKeyframes |= trackMask
				}

				if cb.seenKeyframes&trackMask != 0 {
					if !cb.skipCluster || timecode != cast.time.sent {
						cb.skipCluster = cb.write(cluster)
					}
					if !cb.skipCluster || !cb.write(buf) {
						cb.seenKeyframes &= ^trackMask
					}
				}
			}

			cast.time.sent = timecode

		default:
			return 0, errors.New("unknown EBML tag")
		}

		cast.buffer = cast.buffer[len(buf):]
	}
}
<commit_msg>Hold a mutex while adding/removing clients.<commit_after>package main

import (
	"errors"
	"sync"
)

type broadcastViewer struct {
	// This function may return `false` to signal that it cannot write any more data.
	// The stream will resynchronize at next keyframe.
	write func(data []byte) bool
	// Viewers may hop between streams, but should only receive headers once.
	// This includes track info, as codecs must stay the same between segments.
	skipHeaders bool
	// We group blocks into indeterminate-length clusters. So long as
	// the cluster's timecode has not changed, there's no need to start a new one.
	skipCluster bool
	// To avoid decoding errors due to missing reference frames, the first
	// frame of each track received by a viewer must be a keyframe.
	// Each track for which a keyframe has been sent is marked by a bit here.
	seenKeyframes uint32
}

type Broadcast struct {
	Closed   bool // When set to `true`, all viewers receive an empty bytearray as a notification.
	HasVideo bool
	HasAudio bool
	Width    uint // Dimensions of the video track that came last in the `Tracks` tag.
	Height   uint // Hopefully, there's only one video track in the file.

	vlock   sync.Mutex // protects `viewers`. not RWMutex because there's only one reader.
	viewers map[chan<- []byte]*broadcastViewer
	buffer  []byte
	header  []byte // The EBML (DocType) tag.
	tracks  []byte // The beginning of the Segment (Tracks + Info).

	time struct {
		last  uint64 // Last seen block timecode. The next timecode must be no less than that.
		recv  uint64 // Last received cluster timecode, shifted to ensure monotonicity.
		sent  uint64 // Last sent cluster timecode. (All viewers receive same clusters.)
		shift uint64 // By how much the cluster timecode has been shifted.
	}
}

func NewBroadcast() Broadcast {
	return Broadcast{viewers: make(map[chan<- []byte]*broadcastViewer)}
}

func (cast *Broadcast) Close() error {
	cast.Closed = true
	for _, cb := range cast.viewers {
		cb.write([]byte{})
	}
	return nil
}

func (cast *Broadcast) Connect(ch chan<- []byte, skipHeaders bool) {
	write := func(data []byte) bool {
		// `Broadcast.Write` emits data in block-sized chunks.
		// Thus the buffer size is measured in frames, not bytes.
		if len(ch) == cap(ch) {
			return false
		}
		ch <- data
		return true
	}

	cast.vlock.Lock()
	cast.viewers[ch] = &broadcastViewer{write, skipHeaders, false, 0}
	cast.vlock.Unlock()
}

func (cast *Broadcast) Disconnect(ch chan<- []byte) {
	cast.vlock.Lock()
	delete(cast.viewers, ch)
	cast.vlock.Unlock()
}

func (cast *Broadcast) Reset() {
	cast.buffer = nil
}

func (cast *Broadcast) Write(data []byte) (int, error) {
	cast.buffer = append(cast.buffer, data...)

	for {
		buf := cast.buffer
		tag := EBMLParseTagIncomplete(buf)
		if tag.Consumed == 0 {
			return len(data), nil
		}

		if tag.ID == EBMLSegmentTag || tag.ID == EBMLTracksTag || tag.ID == EBMLClusterTag {
			// Parse the contents of these tags in the same loop.
			buf = buf[:tag.Consumed]
			// Chrome crashes if an indeterminate length is not encoded as 0xFF.
			// If we want to recode it, we'll also need some space for a Void tag.
			if tag.Length == EBMLIndeterminate && tag.Consumed >= 7 {
				cast.buffer[4] = 0xFF
				cast.buffer[5] = EBMLVoidTag
				cast.buffer[6] = 0x80 | byte(tag.Consumed-7)
			}
		} else {
			total := tag.Length + uint64(tag.Consumed)
			if total > 1024*1024 {
				return 0, errors.New("data block too big")
			}

			if total > uint64(len(buf)) {
				return len(data), nil
			}

			buf = buf[:total]
		}

		switch tag.ID {
		case EBMLSeekHeadTag:
			// Disallow seeking.
		case EBMLChaptersTag:
			// Disallow seeking again.
		case EBMLCuesTag:
			// Disallow even more seeking.
		case EBMLVoidTag:
			// Waste of space.
		case EBMLTagsTag:
			// Maybe later.
		case EBMLClusterTag:
			// Ignore boundaries, we'll regroup the data anyway.
		case EBMLPrevSizeTag:
			// Disallow backward seeking too.

		case EBMLEBMLTag:
			// The header is the same in all WebM-s.
			if len(cast.header) == 0 {
				cast.header = append([]byte{}, buf...)
			}

		case EBMLSegmentTag:
			cast.HasVideo = false
			cast.HasAudio = false
			cast.Width = 0
			cast.Height = 0
			cast.tracks = append([]byte{}, buf...)
			// Will recalculate this when the first block arrives.
			cast.time.shift = 0

		case EBMLInfoTag:
			// Default timecode resolution in Matroska is 1 ms. This value is required
			// in WebM; we'll check just in case. Obviously, our timecode rewriting
			// logic won't work with non-millisecond resolutions.
			var scale uint64 = 0

			for buf2 := tag.Contents(buf); len(buf2) != 0; {
				tag2 := EBMLParseTag(buf2)

				switch tag2.ID {
				case 0:
					return 0, errors.New("malformed EBML")

				case EBMLDurationTag:
					total := tag2.Length + uint64(tag2.Consumed) - 2
					if total > 0x7F {
						// I'd rather avoid shifting memory. What kind of integer
						// needs 128 bytes, anyway?
						return 0, errors.New("EBML Duration too large")
					}
					// Live streams must not have a duration.
					buf2[0] = EBMLVoidTag
					buf2[1] = 0x80 | byte(total)

				case EBMLTimecodeScaleTag:
					scale = EBMLParseFixedUint(tag2.Contents(buf2))
				}

				buf2 = tag2.Skip(buf2)
			}

			if scale != 1000000 {
				return 0, errors.New("invalid timecode scale")
			}

			cast.tracks = append(cast.tracks, buf...)

		case EBMLTrackEntryTag:
			// Since `viewer.seenKeyframes` is a 32-bit vector,
			// we need to check that there are at most 32 tracks.
			for buf2 := tag.Contents(buf); len(buf2) != 0; {
				tag2 := EBMLParseTag(buf2)

				switch tag2.ID {
				case 0:
					return 0, errors.New("malformed EBML")

				case EBMLTrackNumberTag:
					// go needs sizeof.
					if t := EBMLParseFixedUint(tag2.Contents(buf2)); t >= 32 {
						return 0, errors.New("too many tracks?")
					}

				case EBMLAudioTag:
					cast.HasAudio = true

				case EBMLVideoTag:
					cast.HasVideo = true
					// While we're here, let's grab some metadata, too.
					for buf3 := tag2.Contents(buf2); len(buf3) != 0; {
						tag3 := EBMLParseTag(buf3)

						switch tag3.ID {
						case 0:
							return 0, errors.New("malformed EBML")

						case EBMLPixelWidthTag:
							cast.Width = uint(EBMLParseFixedUint(tag3.Contents(buf3)))

						case EBMLPixelHeightTag:
							cast.Height = uint(EBMLParseFixedUint(tag3.Contents(buf3)))
						}

						buf3 = tag3.Skip(buf3)
					}
				}

				buf2 = tag2.Skip(buf2)
			}

			cast.tracks = append(cast.tracks, buf...)

		case EBMLTracksTag:
			cast.tracks = append(cast.tracks, buf...)

		case EBMLTimecodeTag:
			// Will reencode it when sending a Cluster.
			cast.time.recv = EBMLParseFixedUint(tag.Contents(buf)) + cast.time.shift

		case EBMLBlockGroupTag, EBMLSimpleBlockTag:
			key := false
			block := tag.Contents(buf)

			if tag.ID == EBMLBlockGroupTag {
				key, block = true, nil

				for buf2 := tag.Contents(buf); len(buf2) != 0; {
					tag2 := EBMLParseTag(buf2)

					switch tag2.ID {
					case 0:
						return 0, errors.New("malformed EBML")

					case EBMLBlockTag:
						block = tag2.Contents(buf2)

					case EBMLReferenceBlockTag:
						// Keyframes, by definition, have no reference frame.
						key = EBMLParseFixedUint(tag2.Contents(buf2)) == 0
					}

					buf2 = tag2.Skip(buf2)
				}

				if block == nil {
					return 0, errors.New("a BlockGroup contains no Blocks")
				}
			}

			track := EBMLParseUint(block)
			if track.Consumed == 0 || track.Value >= 32 || len(block) < track.Consumed+3 {
				return 0, errors.New("invalid track")
			}
			// This bit is always 0 in a Block, but 1 in a keyframe SimpleBlock.
			key = key || block[track.Consumed+2]&0x80 != 0
			// Block timecodes are relative to cluster ones.
			timecode := uint64(block[track.Consumed+0])<<8 | uint64(block[track.Consumed+1])
			if cast.time.recv+timecode < cast.time.last {
				cast.time.shift += cast.time.last - (cast.time.recv + timecode)
				cast.time.recv = cast.time.last - timecode
			}
			cast.time.last = cast.time.recv + timecode

			ctc := cast.time.recv
			cluster := []byte{
				EBMLClusterTag >> 24 & 0xFF,
				EBMLClusterTag >> 16 & 0xFF,
				EBMLClusterTag >> 8 & 0xFF,
				EBMLClusterTag & 0xFF, 0xFF,
				EBMLTimecodeTag, 0x88,
				byte(ctc >> 56), byte(ctc >> 48), byte(ctc >> 40), byte(ctc >> 32),
				byte(ctc >> 24), byte(ctc >> 16), byte(ctc >> 8), byte(ctc),
			}

			trackMask := uint32(1) << track.Value
			cast.vlock.Lock()
			for _, cb := range cast.viewers {
				if !cb.skipHeaders {
					if !cb.write(cast.header) || !cb.write(cast.tracks) {
						continue
					}

					cb.skipHeaders = true
					cb.skipCluster = false
				}

				if key {
					cb.seenKeyframes |= trackMask
				}

				if cb.seenKeyframes&trackMask != 0 {
					if !cb.skipCluster || timecode != cast.time.sent {
						cb.skipCluster = cb.write(cluster)
					}
					if !cb.skipCluster || !cb.write(buf) {
						cb.seenKeyframes &= ^trackMask
					}
				}
			}

			cast.vlock.Unlock()
			cast.time.sent = timecode

		default:
			return 0, errors.New("unknown EBML tag")
		}

		cast.buffer = cast.buffer[len(buf):]
	}
}
<|endoftext|>


<commit_before>package rpc

import (
	"bufio"
	"bytes"
	"code.google.com/p/goprotobuf/proto"
	"encoding/binary"
	"errors"
	"fmt"
	hdfs "github.com/colinmarc/hdfs/protocol/hadoop_hdfs"
	"hash/crc32"
	"io"
	"io/ioutil"
	"math"
	"net"
)

const (
	dataTransferVersion = 0x1c
	readBlockOp         = 0x51
)

// TODO: datanode blacklisting

// BlockReader implements io.ReaderCloser for reading a single block from HDFS,
// abstracting over reading from multiple datanodes.
type BlockReader struct {
	closed bool
	conn   net.Conn
	reader *bufio.Reader

	block       *hdfs.LocatedBlockProto
	checksumTab *crc32.Table

	offset    uint64
	chunkSize uint32
	packet    openPacket
	buf       bytes.Buffer
}

type openPacket struct {
	numChunks     int
	nextChunk     int
	checksumBytes []byte
	blockOffset   uint64
	packetOffset  uint64
	length        uint64
	last          bool
}

// NewBlockReader returns a new BlockReader, given the block information and
// security token from the namenode.
func NewBlockReader(block *hdfs.LocatedBlockProto, offset uint64) (*BlockReader, error) {
	br := &BlockReader{
		block:  block,
		offset: offset,
	}

	// TODO check multiple datanodes
	datanode := br.block.GetLocs()[0].GetId()
	address := fmt.Sprintf("%s:%d", datanode.GetIpAddr(), datanode.GetXferPort())
	err := br.connect(address)
	if err != nil {
		return nil, err
	}

	return br, nil
}

func (br *BlockReader) connect(datanode string) error {
	conn, err := net.DialTimeout("tcp", datanode, connectionTimeout)
	if err != nil {
		return err
	}

	br.conn = conn
	err = br.writeBlockReadRequest()
	if err != nil {
		return err
	}

	br.reader = bufio.NewReader(br.conn)
	resp, err := br.readBlockReadResponse()
	if err != nil {
		return err
	}

	checksumInfo := resp.GetReadOpChecksumInfo().GetChecksum()
	checksumType := checksumInfo.GetType()
	if checksumType == hdfs.ChecksumTypeProto_CHECKSUM_CRC32 {
		br.checksumTab = crc32.IEEETable
	} else if checksumType == hdfs.ChecksumTypeProto_CHECKSUM_CRC32C {
		br.checksumTab = crc32.MakeTable(crc32.Castagnoli)
	} else {
		return fmt.Errorf("Unsupported checksum type:", checksumType)
	}

	br.chunkSize = checksumInfo.GetBytesPerChecksum()
	br.startNewPacket()

	// The read will start aligned to a chunk boundary, so we need to seek forward
	// to the requested offset.
	amountToDiscard := br.offset - br.packet.blockOffset
	if amountToDiscard > 0 {
		io.CopyN(ioutil.Discard, br, int64(amountToDiscard))
	}

	return nil
}

func (br *BlockReader) Close() {
	br.conn.Close()
	br.closed = true
}

func (br *BlockReader) Read(b []byte) (int, error) {
	if br.closed {
		return 0, io.ErrClosedPipe
	} else if br.offset >= br.block.GetB().GetNumBytes() {
		br.Close()
		return 0, io.EOF
	}

	// first, read any leftover data from buf
	if br.buf.Len() > 0 {
		n, _ := br.buf.Read(b)
		return n, nil
	}

	if br.packet.nextChunk >= br.packet.numChunks {
		if br.packet.last {
			return 0, io.EOF
		}

		br.startNewPacket()
	}

	// then, read until we fill up b or we reach the end of the packet
	readOffset := 0
	for br.packet.nextChunk < br.packet.numChunks {
		chOff := 4 * br.packet.nextChunk
		checksum := br.packet.checksumBytes[chOff : chOff+4]

		remaining := br.packet.length - br.packet.packetOffset
		chunkLength := int64(math.Min(float64(br.chunkSize), float64(remaining)))

		chunkReader := io.LimitReader(br.reader, int64(chunkLength))
		chunkBytes := b[readOffset:]
		n, err := chunkReader.Read(chunkBytes)

		readOffset += n
		br.packet.packetOffset += uint64(n)
		br.packet.nextChunk++

		if err != nil {
			if err == io.EOF {
				err = io.ErrUnexpectedEOF
			}

			br.Close()
			return readOffset, err
		}

		crc := crc32.Checksum(chunkBytes[:n], br.checksumTab)

		if int64(n) < chunkLength {
			// save any leftovers
			br.buf.Reset()
			leftover, err := br.buf.ReadFrom(chunkReader)
			if err != nil {
				return readOffset, err
			}

			br.packet.packetOffset += uint64(leftover)

			// update the checksum with the leftovers
			crc = crc32.Update(crc, br.checksumTab, br.buf.Bytes())
		}

		if crc != binary.BigEndian.Uint32(checksum) {
			return readOffset, errors.New("Invalid checksum from the datanode!")
		}

		if readOffset == len(b) {
			break
		}
	}

	return readOffset, nil
}

// A read request to a datanode:
// +-----------------------------------------------------------+
// |  Data Transfer Protocol Version, int16                    |
// +-----------------------------------------------------------+
// |  Op code, 1 byte (READ_BLOCK = 0x51)                      |
// +-----------------------------------------------------------+
// |  varint length + OpReadBlockProto                         |
// +-----------------------------------------------------------+
func (br *BlockReader) writeBlockReadRequest() error {
	header := []byte{0x00, dataTransferVersion, readBlockOp}

	// TODO offset/length?
	needed := (br.block.GetB().GetNumBytes() - br.offset)
	op := newReadBlockOp(br.block, br.offset, needed)
	opBytes, err := makeDelimitedMsg(op)
	if err != nil {
		return err
	}

	req := append(header, opBytes...)
	_, err = br.conn.Write(req)
	if err != nil {
		return err
	}

	return nil
}

// The initial response from the datanode:
// +-----------------------------------------------------------+
// |  varint length + BlockOpResponseProto                     |
// +-----------------------------------------------------------+
func (br *BlockReader) readBlockReadResponse() (*hdfs.BlockOpResponseProto, error) {
	respLength, err := binary.ReadUvarint(br.reader)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}

		return nil, err
	}

	respBytes := make([]byte, respLength)
	_, err = io.ReadFull(br.reader, respBytes)
	if err != nil {
		return nil, err
	}

	resp := &hdfs.BlockOpResponseProto{}
	err = proto.Unmarshal(respBytes, resp)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

// A packet from the datanode:
// +-----------------------------------------------------------+
// |  uint32 length of the packet                              |
// +-----------------------------------------------------------+
// |  size of the PacketHeaderProto, uint16                    |
// +-----------------------------------------------------------+
// |  PacketHeaderProto                                        |
// +-----------------------------------------------------------+
// |  N checksums, 4 bytes each                                |
// +-----------------------------------------------------------+
// |  N chunks of payload data                                 |
// +-----------------------------------------------------------+
func (br *BlockReader) startNewPacket() error {
	header, err := br.readPacketHeader()
	if err != nil {
		return err
	}

	blockOffset := uint64(header.GetOffsetInBlock())
	dataLength := uint64(header.GetDataLen())
	numChunks := int(math.Ceil(float64(dataLength) / float64(br.chunkSize)))

	// TODO don't assume checksum size is 4
	br.packet = openPacket{
		numChunks:     numChunks,
		nextChunk:     0,
		checksumBytes: make([]byte, numChunks*4),
		blockOffset:   blockOffset,
		packetOffset:  0,
		length:        dataLength,
		last:          header.GetLastPacketInBlock(),
	}

	_, err = io.ReadFull(br.reader, br.packet.checksumBytes)
	if err != nil {
		return err
	}

	return nil
}

func (br *BlockReader) readPacketHeader() (*hdfs.PacketHeaderProto, error) {
	var packetLength uint32
	err := binary.Read(br.reader, binary.BigEndian, &packetLength)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}

		return nil, err
	}

	var packetHeaderLength uint16
	err = binary.Read(br.reader, binary.BigEndian, &packetHeaderLength)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}

		return nil, err
	}

	packetHeaderBytes := make([]byte, packetHeaderLength)
	_, err = io.ReadFull(br.reader, packetHeaderBytes)
	if err != nil {
		return nil, err
	}

	packetHeader := &hdfs.PacketHeaderProto{}
	err = proto.Unmarshal(packetHeaderBytes, packetHeader)

	return packetHeader, nil
}

func newReadBlockOp(block *hdfs.LocatedBlockProto, offset, length uint64) *hdfs.OpReadBlockProto {
	return &hdfs.OpReadBlockProto{
		Header: &hdfs.ClientOperationHeaderProto{
			BaseHeader: &hdfs.BaseHeaderProto{
				Block: block.GetB(),
				Token: block.GetBlockToken(),
			},
			ClientName: proto.String(ClientName),
		},
		Offset: proto.Uint64(offset),
		Len:    proto.Uint64(length),
	}
}
<commit_msg>fix a bug causing inconsistent reads<commit_after>package rpc

import (
	"bufio"
	"bytes"
	"code.google.com/p/goprotobuf/proto"
	"encoding/binary"
	"errors"
	"fmt"
	hdfs "github.com/colinmarc/hdfs/protocol/hadoop_hdfs"
	"hash/crc32"
	"io"
	"io/ioutil"
	"math"
	"net"
)

const (
	dataTransferVersion = 0x1c
	readBlockOp         = 0x51
)

// TODO: datanode blacklisting

// BlockReader implements io.ReaderCloser for reading a single block from HDFS,
// abstracting over reading from multiple datanodes.
type BlockReader struct {
	closed bool
	conn   net.Conn
	reader *bufio.Reader

	block       *hdfs.LocatedBlockProto
	checksumTab *crc32.Table

	offset    uint64
	chunkSize uint32
	packet    openPacket
	buf       bytes.Buffer
}

type openPacket struct {
	numChunks     int
	nextChunk     int
	checksumBytes []byte
	blockOffset   uint64
	packetOffset  uint64
	length        uint64
	last          bool
}

// NewBlockReader returns a new BlockReader, given the block information and
// security token from the namenode.
func NewBlockReader(block *hdfs.LocatedBlockProto, offset uint64) (*BlockReader, error) {
	br := &BlockReader{
		block:  block,
		offset: offset,
	}

	// TODO check multiple datanodes
	datanode := br.block.GetLocs()[0].GetId()
	address := fmt.Sprintf("%s:%d", datanode.GetIpAddr(), datanode.GetXferPort())
	err := br.connect(address)
	if err != nil {
		return nil, err
	}

	return br, nil
}

func (br *BlockReader) connect(datanode string) error {
	conn, err := net.DialTimeout("tcp", datanode, connectionTimeout)
	if err != nil {
		return err
	}

	br.conn = conn
	err = br.writeBlockReadRequest()
	if err != nil {
		return err
	}

	br.reader = bufio.NewReader(br.conn)
	resp, err := br.readBlockReadResponse()
	if err != nil {
		return err
	}

	checksumInfo := resp.GetReadOpChecksumInfo().GetChecksum()
	checksumType := checksumInfo.GetType()
	if checksumType == hdfs.ChecksumTypeProto_CHECKSUM_CRC32 {
		br.checksumTab = crc32.IEEETable
	} else if checksumType == hdfs.ChecksumTypeProto_CHECKSUM_CRC32C {
		br.checksumTab = crc32.MakeTable(crc32.Castagnoli)
	} else {
		return fmt.Errorf("Unsupported checksum type:", checksumType)
	}

	br.chunkSize = checksumInfo.GetBytesPerChecksum()
	br.startNewPacket()

	// The read will start aligned to a chunk boundary, so we need to seek forward
	// to the requested offset.
	amountToDiscard := br.offset - br.packet.blockOffset
	if amountToDiscard > 0 {
		io.CopyN(ioutil.Discard, br, int64(amountToDiscard))
	}

	return nil
}

func (br *BlockReader) Close() {
	br.conn.Close()
	br.closed = true
}

func (br *BlockReader) Read(b []byte) (int, error) {
	if br.closed {
		return 0, io.ErrClosedPipe
	} else if br.offset >= br.block.GetB().GetNumBytes() {
		br.Close()
		return 0, io.EOF
	}

	// first, read any leftover data from buf
	if br.buf.Len() > 0 {
		n, _ := br.buf.Read(b)
		return n, nil
	}

	if br.packet.nextChunk >= br.packet.numChunks {
		if br.packet.last {
			return 0, io.EOF
		}

		br.startNewPacket()
	}

	// then, read until we fill up b or we reach the end of the packet
	readOffset := 0
	for br.packet.nextChunk < br.packet.numChunks {
		chOff := 4 * br.packet.nextChunk
		checksum := br.packet.checksumBytes[chOff : chOff+4]

		remaining := br.packet.length - br.packet.packetOffset
		chunkLength := int64(math.Min(float64(br.chunkSize), float64(remaining)))

		chunkReader := io.LimitReader(br.reader, int64(chunkLength))
		chunkBytes := b[readOffset:]
		bytesToRead := int(math.Min(float64(len(chunkBytes)), float64(chunkLength)))
		n, err := io.ReadAtLeast(chunkReader, chunkBytes, bytesToRead)

		readOffset += n
		br.packet.packetOffset += uint64(n)
		br.packet.nextChunk++

		if err != nil {
			br.Close()
			return readOffset, err
		}

		crc := crc32.Checksum(chunkBytes[:n], br.checksumTab)

		if int64(n) < chunkLength {
			// save any leftovers
			br.buf.Reset()
			leftover, err := br.buf.ReadFrom(chunkReader)
			if err != nil {
				return readOffset, err
			}

			br.packet.packetOffset += uint64(leftover)

			// update the checksum with the leftovers
			crc = crc32.Update(crc, br.checksumTab, br.buf.Bytes())
		}

		if crc != binary.BigEndian.Uint32(checksum) {
			return readOffset, errors.New("Invalid checksum from the datanode!")
		}

		if readOffset == len(b) {
			break
		}
	}

	return readOffset, nil
}

// A read request to a datanode:
// +-----------------------------------------------------------+
// |  Data Transfer Protocol Version, int16                    |
// +-----------------------------------------------------------+
// |  Op code, 1 byte (READ_BLOCK = 0x51)                      |
// +-----------------------------------------------------------+
// |  varint length + OpReadBlockProto                         |
// +-----------------------------------------------------------+
func (br *BlockReader) writeBlockReadRequest() error {
	header := []byte{0x00, dataTransferVersion, readBlockOp}

	// TODO offset/length?
	needed := (br.block.GetB().GetNumBytes() - br.offset)
	op := newReadBlockOp(br.block, br.offset, needed)
	opBytes, err := makeDelimitedMsg(op)
	if err != nil {
		return err
	}

	req := append(header, opBytes...)
	_, err = br.conn.Write(req)
	if err != nil {
		return err
	}

	return nil
}

// The initial response from the datanode:
// +-----------------------------------------------------------+
// |  varint length + BlockOpResponseProto                     |
// +-----------------------------------------------------------+
func (br *BlockReader) readBlockReadResponse() (*hdfs.BlockOpResponseProto, error) {
	respLength, err := binary.ReadUvarint(br.reader)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}

		return nil, err
	}

	respBytes := make([]byte, respLength)
	_, err = io.ReadFull(br.reader, respBytes)
	if err != nil {
		return nil, err
	}

	resp := &hdfs.BlockOpResponseProto{}
	err = proto.Unmarshal(respBytes, resp)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

// A packet from the datanode:
// +-----------------------------------------------------------+
// |  uint32 length of the packet                              |
// +-----------------------------------------------------------+
// |  size of the PacketHeaderProto, uint16                    |
// +-----------------------------------------------------------+
// |  PacketHeaderProto                                        |
// +-----------------------------------------------------------+
// |  N checksums, 4 bytes each                                |
// +-----------------------------------------------------------+
// |  N chunks of payload data                                 |
// +-----------------------------------------------------------+
func (br *BlockReader) startNewPacket() error {
	header, err := br.readPacketHeader()
	if err != nil {
		return err
	}

	blockOffset := uint64(header.GetOffsetInBlock())
	dataLength := uint64(header.GetDataLen())
	numChunks := int(math.Ceil(float64(dataLength) / float64(br.chunkSize)))

	// TODO don't assume checksum size is 4
	br.packet = openPacket{
		numChunks:     numChunks,
		nextChunk:     0,
		checksumBytes: make([]byte, numChunks*4),
		blockOffset:   blockOffset,
		packetOffset:  0,
		length:        dataLength,
		last:          header.GetLastPacketInBlock(),
	}

	_, err = io.ReadFull(br.reader, br.packet.checksumBytes)
	if err != nil {
		return err
	}

	return nil
}

func (br *BlockReader) readPacketHeader() (*hdfs.PacketHeaderProto, error) {
	var packetLength uint32
	err := binary.Read(br.reader, binary.BigEndian, &packetLength)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}

		return nil, err
	}

	var packetHeaderLength uint16
	err = binary.Read(br.reader, binary.BigEndian, &packetHeaderLength)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}

		return nil, err
	}

	packetHeaderBytes := make([]byte, packetHeaderLength)
	_, err = io.ReadFull(br.reader, packetHeaderBytes)
	if err != nil {
		return nil, err
	}

	packetHeader := &hdfs.PacketHeaderProto{}
	err = proto.Unmarshal(packetHeaderBytes, packetHeader)

	return packetHeader, nil
}

func newReadBlockOp(block *hdfs.LocatedBlockProto, offset, length uint64) *hdfs.OpReadBlockProto {
	return &hdfs.OpReadBlockProto{
		Header: &hdfs.ClientOperationHeaderProto{
			BaseHeader: &hdfs.BaseHeaderProto{
				Block: block.GetB(),
				Token: block.GetBlockToken(),
			},
			ClientName: proto.String(ClientName),
		},
		Offset: proto.Uint64(offset),
		Len:    proto.Uint64(length),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

type Archive struct {
	*ArchiveResult
	Error string `json:"error,omitempty"`
}

type Site struct {
	Href     string `json:"href"`
	Coords   string `json:"coords"`
	Title    string `json:"title"`
	Response struct {
		Status     int    `json:"status,omitempty"`
		Error      string `json:"error,omitempty"`
		Size       int64  `json:"size,omitempty"`
		Title      string `json:"title,omitempty"`
		Redirected string `json:"redirected,omitempty"`
		Squatter   bool   `json:"squatter,omitempty"`
	} `json:"response,omitempty"`
	Archive *Archive `json:"archive,omitempty"`
}

// Records returns a CSV-compatible flat output
func (s *Site) Records() []string {
	return []string{
		s.Href, s.Coords, s.Title,
		fmt.Sprintf("%d", s.Response.Status),
		s.Response.Error,
		fmt.Sprintf("%d", s.Response.Size),
		s.Response.Title,
		s.Response.Redirected,
		fmt.Sprintf("%t", s.Response.Squatter),
	}
}

var skip = flag.Int("skip", 0, "skip this many lines before processing")
var limit = flag.Int("limit", -1, "abort after processing this many lines")
var concurrency = flag.Int("c", 1, "concurrency")
var exportCSV = flag.Bool("csv", false, "export in csv format")
var loadArchive = flag.Bool("archive", false, "load snapshot from archive.org")

const maxReadSize = 262144

func main() {
	flag.Parse()

	// Read a JSON per line
	var dec = json.NewDecoder(os.Stdin)
	var enc Encoder

	if *exportCSV {
		enc = &CSVEncoder{csv.NewWriter(os.Stdout)}
	} else {
		enc = json.NewEncoder(os.Stdout)
	}

	client := &http.Client{
		Timeout: 10 * time.Second,
	}

	if *concurrency < 1 {
		*concurrency = 1
	}
	if *skip < 0 {
		*skip = 0
	}

	archiveMutex := sync.Mutex{}
	workers := make(chan worker, *concurrency)
	for i := 0; i < *concurrency; i++ {
		w := worker{
			client: client,
		}
		if *loadArchive {
			w.loadArchive = true
			w.lockArchive = &archiveMutex
		}
		workers <- w
	}
	results := make(chan Site)
	wait := sync.WaitGroup{}

	go func() {
		// Writer
		for s := range results {
			if err := enc.Encode(&s); err != nil {
				log.Fatalf("encode failed: %s", err)
			}
			wait.Done()
		}
	}()

	var i int
	var s Site
	for {
		if *limit > 0 && *limit <= i-*skip {
			log.Printf("stopping early due to limit: %d (limit=%d, skip=%d)", i, *limit, *skip)
			break
		}

		if err := dec.Decode(&s); err == io.EOF {
			break
		} else if err != nil {
			log.Fatalf("decode failed on line %d: %s", i, err)
		}

		if i < *skip {
			continue
		}

		w := <-workers
		go func(s Site) {
			wait.Add(1)
			log.Printf("[%d] Processing %s", i, s.Href)

			results <- w.Process(s)
			workers <- w
		}(s)

		i++
	}

	log.Printf("shutting down after line %d", i)
	wait.Wait()
	close(results)
	close(workers)
}

type worker struct {
	client      *http.Client
	loadArchive bool
	lockArchive sync.Locker
}

// Process takes a Site query and returns an augmented Site
func (w *worker) Process(s Site) Site {
	if !w.loadArchive {
		w.lockArchive.Lock()

		archive := ArchiveClient{w.client}
		// The auction ended in Jan 2006, so we fetch archives after Feb
		results, err := archive.History(s.Href, Param{"limit", "1"}, Param{"from", "200602"})
		arch := &Archive{}
		if err != nil {
			arch.Error = err.Error()
			log.Printf("archive error: %s", err)
		} else if len(results) > 0 {
			arch.ArchiveResult = &results[0]
		} else {
			arch = nil
		}
		s.Archive = arch
		w.lockArchive.Unlock()
	}

	resp, err := w.client.Get(s.Href)
	if err != nil {
		s.Response.Error = err.Error()
		return s
	}
	defer resp.Body.Close()

	// Fill in the response
	s.Response.Status = resp.StatusCode
	if resp.ContentLength >= 0 {
		s.Response.Size = resp.ContentLength
	}

	// Is it a redirect?
	url := resp.Request.URL.String()
	if url != s.Href {
		s.Response.Redirected = url
	}

	// What's in the body?
	readcount := &CountingReader{Reader: resp.Body}
	parsed := ParseHTML(readcount)
	s.Response.Title = parsed.Title
	s.Response.Squatter = parsed.MentionsDomain

	// Should we try to guess the content size?
	if s.Response.Redirected == "" && s.Response.Size == 0 {
		// Consume the body until completion to measure the body
		for {
			if _, err := readcount.Read(nil); err == io.EOF {
				break
			}
			if readcount.Count() > maxReadSize {
				break
			}
		}
		s.Response.Size = int64(readcount.Count())
	}

	return s
}
<commit_msg>crawler: fix archive inverse if<commit_after>package main

import (
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

type Archive struct {
	*ArchiveResult
	Error string `json:"error,omitempty"`
}

type Site struct {
	Href     string `json:"href"`
	Coords   string `json:"coords"`
	Title    string `json:"title"`
	Response struct {
		Status     int    `json:"status,omitempty"`
		Error      string `json:"error,omitempty"`
		Size       int64  `json:"size,omitempty"`
		Title      string `json:"title,omitempty"`
		Redirected string `json:"redirected,omitempty"`
		Squatter   bool   `json:"squatter,omitempty"`
	} `json:"response,omitempty"`
	Archive *Archive `json:"archive,omitempty"`
}

// Records returns a CSV-compatible flat output
func (s *Site) Records() []string {
	return []string{
		s.Href, s.Coords, s.Title,
		fmt.Sprintf("%d", s.Response.Status),
		s.Response.Error,
		fmt.Sprintf("%d", s.Response.Size),
		s.Response.Title,
		s.Response.Redirected,
		fmt.Sprintf("%t", s.Response.Squatter),
	}
}

var skip = flag.Int("skip", 0, "skip this many lines before processing")
var limit = flag.Int("limit", -1, "abort after processing this many lines")
var concurrency = flag.Int("c", 1, "concurrency")
var exportCSV = flag.Bool("csv", false, "export in csv format")
var loadArchive = flag.Bool("archive", false, "load snapshot from archive.org")

const maxReadSize = 262144

func main() {
	flag.Parse()

	// Read a JSON per line
	var dec = json.NewDecoder(os.Stdin)
	var enc Encoder

	if *exportCSV {
		enc = &CSVEncoder{csv.NewWriter(os.Stdout)}
	} else {
		enc = json.NewEncoder(os.Stdout)
	}

	client := &http.Client{
		Timeout: 10 * time.Second,
	}

	if *concurrency < 1 {
		*concurrency = 1
	}
	if *skip < 0 {
		*skip = 0
	}

	archiveMutex := sync.Mutex{}
	workers := make(chan worker, *concurrency)
	for i := 0; i < *concurrency; i++ {
		w := worker{
			client: client,
		}
		if *loadArchive {
			w.loadArchive = true
			w.lockArchive = &archiveMutex
		}
		workers <- w
	}
	results := make(chan Site)
	wait := sync.WaitGroup{}

	go func() {
		// Writer
		for s := range results {
			if err := enc.Encode(&s); err != nil {
				log.Fatalf("encode failed: %s", err)
			}
			wait.Done()
		}
	}()

	var i int
	var s Site
	for {
		if *limit > 0 && *limit <= i-*skip {
			log.Printf("stopping early due to limit: %d (limit=%d, skip=%d)", i, *limit, *skip)
			break
		}

		if err := dec.Decode(&s); err == io.EOF {
			break
		} else if err != nil {
			log.Fatalf("decode failed on line %d: %s", i, err)
		}

		if i < *skip {
			continue
		}

		w := <-workers
		go func(s Site) {
			wait.Add(1)
			log.Printf("[%d] Processing %s", i, s.Href)

			results <- w.Process(s)
			workers <- w
		}(s)

		i++
	}

	log.Printf("shutting down after line %d", i)
	wait.Wait()
	close(results)
	close(workers)
}

type worker struct {
	client      *http.Client
	loadArchive bool
	lockArchive sync.Locker
}

// Process takes a Site query and returns an augmented Site
func (w *worker) Process(s Site) Site {
	if w.loadArchive {
		w.lockArchive.Lock()

		archive := ArchiveClient{w.client}
		// The auction ended in Jan 2006, so we fetch archives after Feb
		results, err := archive.History(s.Href, Param{"limit", "1"}, Param{"from", "200602"})
		arch := &Archive{}
		if err != nil {
			arch.Error = err.Error()
			log.Printf("archive error: %s", err)
		} else if len(results) > 0 {
			arch.ArchiveResult = &results[0]
		} else {
			arch = nil
		}
		s.Archive = arch
		w.lockArchive.Unlock()
	}

	resp, err := w.client.Get(s.Href)
	if err != nil {
		s.Response.Error = err.Error()
		return s
	}
	defer resp.Body.Close()

	// Fill in the response
	s.Response.Status = resp.StatusCode
	if resp.ContentLength >= 0 {
		s.Response.Size = resp.ContentLength
	}

	// Is it a redirect?
	url := resp.Request.URL.String()
	if url != s.Href {
		s.Response.Redirected = url
	}

	// What's in the body?
	readcount := &CountingReader{Reader: resp.Body}
	parsed := ParseHTML(readcount)
	s.Response.Title = parsed.Title
	s.Response.Squatter = parsed.MentionsDomain

	// Should we try to guess the content size?
	if s.Response.Redirected == "" && s.Response.Size == 0 {
		// Consume the body until completion to measure the body
		for {
			if _, err := readcount.Read(nil); err == io.EOF {
				break
			}
			if readcount.Count() > maxReadSize {
				break
			}
		}
		s.Response.Size = int64(readcount.Count())
	}

	return s
}
<|endoftext|>


<commit_before>package dirx

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/json-iterator/go"

	"github.com/qiniu/log"

	"github.com/qiniu/logkit/conf"
	"github.com/qiniu/logkit/reader"
	. "github.com/qiniu/logkit/utils/models"
)

var (
	_ reader.DaemonReader = &Reader{}
	_ reader.StatsReader  = &Reader{}
	_ reader.Reader       = &Reader{}
	_ Resetable           = &Reader{}
)

func init() {
	reader.RegisterConstructor(reader.ModeDirx, NewReader)
}

type message struct {
	result  string
	logpath string
}

type Reader struct {
	meta *reader.Meta
	// Note: 原子操作，用于表示 reader 整体的运行状态
	status int32

	stopChan chan struct{}
	msgChan  chan message
	errChan  chan error

	stats     StatsInfo
	statsLock sync.RWMutex

	headRegexp  *regexp.Regexp
	currentFile string
	dirReaders  *dirReaders

	// 以下为传入参数
	logPathPattern     string
	statInterval       time.Duration
	expire             time.Duration
	submetaExpire      time.Duration
	maxOpenFiles       int
	ignoreHidden       bool
	skipFirstLine      bool
	newFileNewLine     bool
	ignoreFileSuffixes []string
	validFilesRegex    string
	whence             string
	bufferSize         int
}

func NewReader(meta *reader.Meta, conf conf.MapConf) (reader.Reader, error) {
	logPathPattern, err := conf.GetString(reader.KeyLogPath)
	if err != nil {
		return nil, err
	}

	statIntervalDur, _ := conf.GetStringOr(reader.KeyStatInterval, "3m")
	statInterval, err := time.ParseDuration(statIntervalDur)
	if err != nil {
		return nil, err
	}

	expireDur, _ := conf.GetStringOr(reader.KeyExpire, "0s")
	expire, err := time.ParseDuration(expireDur)
	if err != nil {
		return nil, err
	}

	submetaExpireDur, _ := conf.GetStringOr(reader.KeySubmetaExpire, "720h")
	submetaExpire, err := time.ParseDuration(submetaExpireDur)
	if err != nil {
		return nil, err
	}
	// submetaExpire 为 0 时，不清理元数据
	if IsSubmetaExpireValid(submetaExpire, expire) {
		return nil, fmt.Errorf("%q valus is less than %q", reader.KeySubmetaExpire, reader.KeyExpire)
	}

	maxOpenFiles, _ := conf.GetIntOr(reader.KeyMaxOpenFiles, 256)
	ignoreHidden, _ := conf.GetBoolOr(reader.KeyIgnoreHiddenFile, true) // 默认不读取隐藏文件
	skipFirstLine, _ := conf.GetBoolOr(reader.KeySkipFileFirstLine, false)
	newFileNewLine, _ := conf.GetBoolOr(reader.KeyNewFileNewLine, false)
	ignoreFileSuffixes, _ := conf.GetStringListOr(reader.KeyIgnoreFileSuffix, reader.DefaultIgnoreFileSuffixes)
	validFilesRegex, _ := conf.GetStringOr(reader.KeyValidFilePattern, "*")
	whence, _ := conf.GetStringOr(reader.KeyWhence, reader.WhenceOldest)
	bufferSize, _ := conf.GetIntOr(reader.KeyBufSize, reader.DefaultBufSize)

	_, _, bufsize, err := meta.ReadBufMeta()
	if err != nil {
		if os.IsNotExist(err) {
			log.Debugf("Runner[%v] %v recover from meta error %v, ignore...", meta.RunnerName, logPathPattern, err)
		} else {
			log.Warnf("Runner[%v] %v recover from meta error %v, ignore...", meta.RunnerName, logPathPattern, err)
		}
		bufsize = 0
	}

	cachedLines := make(map[string]string)
	buf := make([]byte, bufsize)
	if bufsize > 0 {
		if _, err = meta.ReadBuf(buf); err != nil {
			if os.IsNotExist(err) {
				log.Debugf("Runner[%v] read buf file[%v] failed: %v, ignore...", meta.RunnerName, meta.BufFile(), err)
			} else {
				log.Warnf("Runner[%v] read buf file[%v] failed: %v, ignore...", meta.RunnerName, meta.BufFile(), err)
			}
		} else {
			err = jsoniter.Unmarshal(buf, &cachedLines)
			if err != nil {
				log.Warnf("Runner[%v] unmarshal read buf cache failed: %v, ignore...", meta.RunnerName, err)
			}
		}
	}

	return &Reader{
		meta:               meta,
		status:             reader.StatusInit,
		stopChan:           make(chan struct{}),
		msgChan:            make(chan message),
		errChan:            make(chan error),
		dirReaders:         newDirReaders(meta, expire, cachedLines),
		logPathPattern:     strings.TrimSuffix(logPathPattern, "/"),
		statInterval:       statInterval,
		expire:             expire,
		submetaExpire:      submetaExpire,
		maxOpenFiles:       maxOpenFiles,
		ignoreHidden:       ignoreHidden,
		skipFirstLine:      skipFirstLine,
		newFileNewLine:     newFileNewLine,
		ignoreFileSuffixes: ignoreFileSuffixes,
		validFilesRegex:    validFilesRegex,
		whence:             whence,
		bufferSize:         bufferSize,
	}, nil
}

func (r *Reader) isStopping() bool {
	return atomic.LoadInt32(&r.status) == reader.StatusStopping
}

func (r *Reader) hasStopped() bool {
	return atomic.LoadInt32(&r.status) == reader.StatusStopped
}

func (r *Reader) Name() string {
	return "DirxReader: " + r.logPathPattern
}

func (r *Reader) SetMode(mode string, v interface{}) error {
	reg, err := reader.HeadPatternMode(mode, v)
	if err != nil {
		return fmt.Errorf("get head pattern mode: %v", err)
	}
	if reg != nil {
		r.headRegexp = reg
	}
	return nil
}

func (r *Reader) setStatsError(err string) {
	r.statsLock.Lock()
	defer r.statsLock.Unlock()
	r.stats.LastError = err
}

func (r *Reader) sendError(err error) {
	if err == nil {
		return
	}
	defer func() {
		if rec := recover(); rec != nil {
			log.Errorf("Reader %q was panicked and recovered from %v", r.Name(), rec)
		}
	}()
	r.errChan <- err
}

func (r *Reader) statLogPath() {
	// 达到最大打开文件数时不再追踪
	if r.dirReaders.Num() >= r.maxOpenFiles {
		log.Warnf("Runner[%v] has met 'maxOpenFiles' limit %v, stat new log ignored", r.meta.RunnerName, r.maxOpenFiles)
		return
	}

	matches, err := filepath.Glob(r.logPathPattern)
	if err != nil {
		errMsg := fmt.Sprintf("Runner[%v] stat log path failed: %v", r.meta.RunnerName, err)
		log.Error(errMsg)
		r.setStatsError(errMsg)
		return
	}
	if len(matches) == 0 {
		log.Debugf("Runner[%v] no match found after stated log path %q", r.meta.RunnerName, r.logPathPattern)
		return
	}
	log.Debugf("Runner[%v] %d matches found after stated log path %q: %v", r.meta.RunnerName, len(matches), r.logPathPattern, matches)

	var newPaths []string
	for _, m := range matches {
		logPath, fi, err := GetRealPath(m)
		if err != nil {
			log.Warnf("Runner[%v] file pattern %v match %v stat failed: %v, ignored this match", r.meta.RunnerName, r.logPathPattern, m, err)
			continue
		}
		if !fi.IsDir() {
			log.Debugf("Runner[%v] %v is a file, mode[dirx] only supports stat directory, ignored this match", r.meta.RunnerName, m)
			continue
		}
		if r.dirReaders.HasReader(logPath) {
			log.Debugf("Runner[%v] %q is collecting, ignored this path", r.meta.RunnerName, logPath)
			continue
		}

		// 过期的文件不追踪，除非之前追踪的并且有日志没读完
		if !r.dirReaders.hasCachedLine(logPath) && HasDirExpired(logPath, r.dirReaders.expire) {
			log.Debugf("Runner[%v] log path %q has expired, ignored this time", r.meta.RunnerName, logPath)
			continue
		}

		dr, err := r.dirReaders.NewReader(newReaderOptions{
			Meta:               r.meta,
			OriginalPath:       m,
			LogPath:            logPath,
			IgnoreHidden:       r.ignoreHidden,
			SkipFirstLine:      r.skipFirstLine,
			NewFileNewLine:     r.newFileNewLine,
			IgnoreFileSuffixes: r.ignoreFileSuffixes,
			ValidFilesRegex:    r.validFilesRegex,
			Whence:             r.whence,
			BufferSize:         r.bufferSize,
			MsgChan:            r.msgChan,
			ErrChan:            r.errChan,
		})
		if err != nil {
			err = fmt.Errorf("create new reader for log path %q failed: %v", logPath, err)
			r.sendError(err)
			log.Errorf("Runner[%v] %v, ignored this path", r.meta.RunnerName, err)
			continue
		}

		if r.headRegexp != nil {
			err = dr.br.SetMode(reader.ReadModeHeadPatternRegexp, r.headRegexp)
			if err != nil {
				errMsg := fmt.Sprintf("Runner[%v] set mode for log path %q failed: %v", r.meta.RunnerName, logPath, err)
				log.Error(errMsg)
				r.setStatsError(errMsg)
			}
		}
		newPaths = append(newPaths, logPath)

		if r.hasStopped() {
			log.Warnf("Runner[%v] created new reader for log path %q but daemon reader has stopped, will not run at this time", r.meta.RunnerName, logPath)
			continue
		}

		go dr.Run()
	}
	if len(newPaths) > 0 {
		log.Infof("Runner[%v] stat log path found %d new log paths: %v", r.meta.RunnerName, len(newPaths), newPaths)
	} else {
		log.Debugf("Runner[%v] stat log path has not found any new log path", r.meta.RunnerName)
	}
}

func (r *Reader) Start() error {
	if r.isStopping() || r.hasStopped() {
		return errors.New("reader is stopping or has stopped")
	} else if !atomic.CompareAndSwapInt32(&r.status, reader.StatusInit, reader.StatusRunning) {
		log.Warnf("Runner[%v] %q daemon has already started and is running", r.meta.RunnerName, r.Name())
		return nil
	}

	go func() {
		ticker := time.NewTicker(r.statInterval)
		defer ticker.Stop()
		for {
			r.dirReaders.checkExpiredDirs()
			r.statLogPath()

			select {
			case <-r.stopChan:
				atomic.StoreInt32(&r.status, reader.StatusStopped)
				log.Infof("Runner[%v] %q daemon has stopped from running", r.meta.RunnerName, r.Name())
				return
			case <-ticker.C:
			}
		}
	}()

	if IsSubMetaExpire(r.submetaExpire, r.expire) {
		go func() {
			ticker := time.NewTicker(time.Hour)
			defer ticker.Stop()
			for {
				r.meta.CheckExpiredSubMetas(r.submetaExpire)

				select {
				case <-r.stopChan:
					return
				case <-ticker.C:
				}
			}
		}()
	}

	log.Infof("Runner[%v] %q daemon has started", r.meta.RunnerName, r.Name())
	return nil
}

func (r *Reader) Source() string {
	return r.currentFile
}

// Note: 对 currentFile 的操作非线程安全，需由上层逻辑保证同步调用 ReadLine
func (r *Reader) ReadLine() (string, error) {
	timer := time.NewTimer(time.Second)
	defer timer.Stop()
	select {
	case msg := <-r.msgChan:
		r.currentFile = msg.logpath
		return msg.result, nil
	case <-timer.C:
		return "", r.readError()
	}

	return "", nil
}

func (r *Reader) Status() StatsInfo {
	r.statsLock.RLock()
	defer r.statsLock.RUnlock()

	for _, dr := range r.dirReaders.getReaders() {
		st := dr.Status()
		if st.LastError != "" {
			r.stats.LastError += "\n<" + dr.originalPath + ">: " + st.LastError
		}
	}
	return r.stats
}

// SyncMeta 从队列取数据时同步队列，作用在于保证数据不重复
func (r *Reader) SyncMeta() {
	data, err := r.dirReaders.SyncMeta()
	if err != nil {
		log.Errorf("Runner[%v] reader %q sync meta failed: %v", r.meta.RunnerName, r.Name(), err)
		return
	}

	err = r.meta.WriteBuf(data, 0, 0, len(data))
	if err != nil {
		log.Errorf("Runner[%v] write meta[%v] failed: %v", r.meta.RunnerName, string(data), err)
		return
	}

	if IsSubMetaExpire(r.submetaExpire, r.expire) {
		r.meta.CleanExpiredSubMetas(r.submetaExpire)
	}
}

func (r *Reader) Close() error {
	if !atomic.CompareAndSwapInt32(&r.status, reader.StatusRunning, reader.StatusStopping) {
		log.Warnf("Runner[%v] reader %q is not running, close operation ignored", r.meta.RunnerName, r.Name())
		return nil
	}
	log.Debugf("Runner[%v] %q daemon is stopping", r.meta.RunnerName, r.Name())
	close(r.stopChan)

	r.dirReaders.Close()
	r.SyncMeta()

	// 在所有 dirReader 关闭完成后再关闭管道
	close(r.msgChan)
	close(r.errChan)
	return nil
}

func (r *Reader) Reset() error {
	return r.dirReaders.Reset()
}

func (r *Reader) readError() error {
	select {
	case err := <-r.errChan:
		return err
	default:
		return nil
	}
}
<commit_msg>fix read dirx panic<commit_after>package dirx

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/json-iterator/go"

	"github.com/qiniu/log"

	"github.com/qiniu/logkit/conf"
	"github.com/qiniu/logkit/reader"
	. "github.com/qiniu/logkit/utils/models"
)

var (
	_ reader.DaemonReader = &Reader{}
	_ reader.StatsReader  = &Reader{}
	_ reader.Reader       = &Reader{}
	_ Resetable           = &Reader{}
)

func init() {
	reader.RegisterConstructor(reader.ModeDirx, NewReader)
}

type message struct {
	result  string
	logpath string
}

type Reader struct {
	meta *reader.Meta
	// Note: 原子操作，用于表示 reader 整体的运行状态
	status int32

	stopChan chan struct{}
	msgChan  chan message
	errChan  chan error

	stats     StatsInfo
	statsLock sync.RWMutex

	headRegexp  *regexp.Regexp
	currentFile string
	dirReaders  *dirReaders

	// 以下为传入参数
	logPathPattern     string
	statInterval       time.Duration
	expire             time.Duration
	submetaExpire      time.Duration
	maxOpenFiles       int
	ignoreHidden       bool
	skipFirstLine      bool
	newFileNewLine     bool
	ignoreFileSuffixes []string
	validFilesRegex    string
	whence             string
	bufferSize         int
}

func NewReader(meta *reader.Meta, conf conf.MapConf) (reader.Reader, error) {
	logPathPattern, err := conf.GetString(reader.KeyLogPath)
	if err != nil {
		return nil, err
	}

	statIntervalDur, _ := conf.GetStringOr(reader.KeyStatInterval, "3m")
	statInterval, err := time.ParseDuration(statIntervalDur)
	if err != nil {
		return nil, err
	}

	expireDur, _ := conf.GetStringOr(reader.KeyExpire, "0s")
	expire, err := time.ParseDuration(expireDur)
	if err != nil {
		return nil, err
	}

	submetaExpireDur, _ := conf.GetStringOr(reader.KeySubmetaExpire, "720h")
	submetaExpire, err := time.ParseDuration(submetaExpireDur)
	if err != nil {
		return nil, err
	}
	// submetaExpire 为 0 时，不清理元数据
	if IsSubmetaExpireValid(submetaExpire, expire) {
		return nil, fmt.Errorf("%q valus is less than %q", reader.KeySubmetaExpire, reader.KeyExpire)
	}

	maxOpenFiles, _ := conf.GetIntOr(reader.KeyMaxOpenFiles, 256)
	ignoreHidden, _ := conf.GetBoolOr(reader.KeyIgnoreHiddenFile, true) // 默认不读取隐藏文件
	skipFirstLine, _ := conf.GetBoolOr(reader.KeySkipFileFirstLine, false)
	newFileNewLine, _ := conf.GetBoolOr(reader.KeyNewFileNewLine, false)
	ignoreFileSuffixes, _ := conf.GetStringListOr(reader.KeyIgnoreFileSuffix, reader.DefaultIgnoreFileSuffixes)
	validFilesRegex, _ := conf.GetStringOr(reader.KeyValidFilePattern, "*")
	whence, _ := conf.GetStringOr(reader.KeyWhence, reader.WhenceOldest)
	bufferSize, _ := conf.GetIntOr(reader.KeyBufSize, reader.DefaultBufSize)

	_, _, bufsize, err := meta.ReadBufMeta()
	if err != nil {
		if os.IsNotExist(err) {
			log.Debugf("Runner[%v] %v recover from meta error %v, ignore...", meta.RunnerName, logPathPattern, err)
		} else {
			log.Warnf("Runner[%v] %v recover from meta error %v, ignore...", meta.RunnerName, logPathPattern, err)
		}
		bufsize = 0
	}

	cachedLines := make(map[string]string)
	buf := make([]byte, bufsize)
	if bufsize > 0 {
		if _, err = meta.ReadBuf(buf); err != nil {
			if os.IsNotExist(err) {
				log.Debugf("Runner[%v] read buf file[%v] failed: %v, ignore...", meta.RunnerName, meta.BufFile(), err)
			} else {
				log.Warnf("Runner[%v] read buf file[%v] failed: %v, ignore...", meta.RunnerName, meta.BufFile(), err)
			}
		} else {
			err = jsoniter.Unmarshal(buf, &cachedLines)
			if err != nil {
				log.Warnf("Runner[%v] unmarshal read buf cache failed: %v, ignore...", meta.RunnerName, err)
			}
		}
	}

	return &Reader{
		meta:               meta,
		status:             reader.StatusInit,
		stopChan:           make(chan struct{}),
		msgChan:            make(chan message),
		errChan:            make(chan error),
		dirReaders:         newDirReaders(meta, expire, cachedLines),
		logPathPattern:     strings.TrimSuffix(logPathPattern, "/"),
		statInterval:       statInterval,
		expire:             expire,
		submetaExpire:      submetaExpire,
		maxOpenFiles:       maxOpenFiles,
		ignoreHidden:       ignoreHidden,
		skipFirstLine:      skipFirstLine,
		newFileNewLine:     newFileNewLine,
		ignoreFileSuffixes: ignoreFileSuffixes,
		validFilesRegex:    validFilesRegex,
		whence:             whence,
		bufferSize:         bufferSize,
	}, nil
}

func (r *Reader) isStopping() bool {
	return atomic.LoadInt32(&r.status) == reader.StatusStopping
}

func (r *Reader) hasStopped() bool {
	return atomic.LoadInt32(&r.status) == reader.StatusStopped
}

func (r *Reader) Name() string {
	return "DirxReader: " + r.logPathPattern
}

func (r *Reader) SetMode(mode string, v interface{}) error {
	reg, err := reader.HeadPatternMode(mode, v)
	if err != nil {
		return fmt.Errorf("get head pattern mode: %v", err)
	}
	if reg != nil {
		r.headRegexp = reg
	}
	return nil
}

func (r *Reader) setStatsError(err string) {
	r.statsLock.Lock()
	defer r.statsLock.Unlock()
	r.stats.LastError = err
}

func (r *Reader) sendError(err error) {
	if err == nil {
		return
	}
	defer func() {
		if rec := recover(); rec != nil {
			log.Errorf("Reader %q was panicked and recovered from %v", r.Name(), rec)
		}
	}()
	r.errChan <- err
}

func (r *Reader) statLogPath() {
	// 达到最大打开文件数时不再追踪
	if r.dirReaders.Num() >= r.maxOpenFiles {
		log.Warnf("Runner[%v] has met 'maxOpenFiles' limit %v, stat new log ignored", r.meta.RunnerName, r.maxOpenFiles)
		return
	}

	matches, err := filepath.Glob(r.logPathPattern)
	if err != nil {
		errMsg := fmt.Sprintf("Runner[%v] stat log path failed: %v", r.meta.RunnerName, err)
		log.Error(errMsg)
		r.setStatsError(errMsg)
		return
	}
	if len(matches) == 0 {
		log.Debugf("Runner[%v] no match found after stated log path %q", r.meta.RunnerName, r.logPathPattern)
		return
	}
	log.Debugf("Runner[%v] %d matches found after stated log path %q: %v", r.meta.RunnerName, len(matches), r.logPathPattern, matches)

	var newPaths []string
	for _, m := range matches {
		logPath, fi, err := GetRealPath(m)
		if err != nil {
			log.Warnf("Runner[%v] file pattern %v match %v stat failed: %v, ignored this match", r.meta.RunnerName, r.logPathPattern, m, err)
			continue
		}
		if !fi.IsDir() {
			log.Debugf("Runner[%v] %v is a file, mode[dirx] only supports stat directory, ignored this match", r.meta.RunnerName, m)
			continue
		}
		if r.dirReaders.HasReader(logPath) {
			log.Debugf("Runner[%v] %q is collecting, ignored this path", r.meta.RunnerName, logPath)
			continue
		}

		// 过期的文件不追踪，除非之前追踪的并且有日志没读完
		if !r.dirReaders.hasCachedLine(logPath) && HasDirExpired(logPath, r.dirReaders.expire) {
			log.Debugf("Runner[%v] log path %q has expired, ignored this time", r.meta.RunnerName, logPath)
			continue
		}

		dr, err := r.dirReaders.NewReader(newReaderOptions{
			Meta:               r.meta,
			OriginalPath:       m,
			LogPath:            logPath,
			IgnoreHidden:       r.ignoreHidden,
			SkipFirstLine:      r.skipFirstLine,
			NewFileNewLine:     r.newFileNewLine,
			IgnoreFileSuffixes: r.ignoreFileSuffixes,
			ValidFilesRegex:    r.validFilesRegex,
			Whence:             r.whence,
			BufferSize:         r.bufferSize,
			MsgChan:            r.msgChan,
			ErrChan:            r.errChan,
		})
		if err != nil {
			err = fmt.Errorf("create new reader for log path %q failed: %v", logPath, err)
			r.sendError(err)
			log.Errorf("Runner[%v] %v, ignored this path", r.meta.RunnerName, err)
			continue
		}

		if r.headRegexp != nil {
			err = dr.br.SetMode(reader.ReadModeHeadPatternRegexp, r.headRegexp)
			if err != nil {
				errMsg := fmt.Sprintf("Runner[%v] set mode for log path %q failed: %v", r.meta.RunnerName, logPath, err)
				log.Error(errMsg)
				r.setStatsError(errMsg)
			}
		}
		newPaths = append(newPaths, logPath)

		if r.hasStopped() || r.isStopping() {
			log.Warnf("Runner[%v] created new reader for log path %q but daemon reader has stopped/is stopping, will not run at this time", r.meta.RunnerName, logPath)
			continue
		}

		go dr.Run()
	}
	if len(newPaths) > 0 {
		log.Infof("Runner[%v] stat log path found %d new log paths: %v", r.meta.RunnerName, len(newPaths), newPaths)
	} else {
		log.Debugf("Runner[%v] stat log path has not found any new log path", r.meta.RunnerName)
	}
}

func (r *Reader) Start() error {
	if r.isStopping() || r.hasStopped() {
		return errors.New("reader is stopping or has stopped")
	} else if !atomic.CompareAndSwapInt32(&r.status, reader.StatusInit, reader.StatusRunning) {
		log.Warnf("Runner[%v] %q daemon has already started and is running", r.meta.RunnerName, r.Name())
		return nil
	}

	go func() {
		ticker := time.NewTicker(r.statInterval)
		defer ticker.Stop()
		for {
			r.dirReaders.checkExpiredDirs()
			r.statLogPath()

			select {
			case <-r.stopChan:
				atomic.StoreInt32(&r.status, reader.StatusStopped)
				log.Infof("Runner[%v] %q daemon has stopped from running", r.meta.RunnerName, r.Name())
				return
			case <-ticker.C:
			}
		}
	}()

	if IsSubMetaExpire(r.submetaExpire, r.expire) {
		go func() {
			ticker := time.NewTicker(time.Hour)
			defer ticker.Stop()
			for {
				r.meta.CheckExpiredSubMetas(r.submetaExpire)

				select {
				case <-r.stopChan:
					return
				case <-ticker.C:
				}
			}
		}()
	}

	log.Infof("Runner[%v] %q daemon has started", r.meta.RunnerName, r.Name())
	return nil
}

func (r *Reader) Source() string {
	return r.currentFile
}

// Note: 对 currentFile 的操作非线程安全，需由上层逻辑保证同步调用 ReadLine
func (r *Reader) ReadLine() (string, error) {
	timer := time.NewTimer(time.Second)
	defer timer.Stop()
	select {
	case msg := <-r.msgChan:
		r.currentFile = msg.logpath
		return msg.result, nil
	case <-timer.C:
		return "", r.readError()
	}

	return "", nil
}

func (r *Reader) Status() StatsInfo {
	r.statsLock.RLock()
	defer r.statsLock.RUnlock()

	for _, dr := range r.dirReaders.getReaders() {
		st := dr.Status()
		if st.LastError != "" {
			r.stats.LastError += "\n<" + dr.originalPath + ">: " + st.LastError
		}
	}
	return r.stats
}

// SyncMeta 从队列取数据时同步队列，作用在于保证数据不重复
func (r *Reader) SyncMeta() {
	data, err := r.dirReaders.SyncMeta()
	if err != nil {
		log.Errorf("Runner[%v] reader %q sync meta failed: %v", r.meta.RunnerName, r.Name(), err)
		return
	}

	err = r.meta.WriteBuf(data, 0, 0, len(data))
	if err != nil {
		log.Errorf("Runner[%v] write meta[%v] failed: %v", r.meta.RunnerName, string(data), err)
		return
	}

	if IsSubMetaExpire(r.submetaExpire, r.expire) {
		r.meta.CleanExpiredSubMetas(r.submetaExpire)
	}
}

func (r *Reader) Close() error {
	if !atomic.CompareAndSwapInt32(&r.status, reader.StatusRunning, reader.StatusStopping) {
		log.Warnf("Runner[%v] reader %q is not running, close operation ignored", r.meta.RunnerName, r.Name())
		return nil
	}
	log.Debugf("Runner[%v] %q daemon is stopping", r.meta.RunnerName, r.Name())
	close(r.stopChan)

	r.dirReaders.Close()
	r.SyncMeta()

	// 在所有 dirReader 关闭完成后再关闭管道
	close(r.msgChan)
	close(r.errChan)
	return nil
}

func (r *Reader) Reset() error {
	return r.dirReaders.Reset()
}

func (r *Reader) readError() error {
	select {
	case err := <-r.errChan:
		return err
	default:
		return nil
	}
}
<|endoftext|>