<commit_before>package multiplex

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	pool "github.com/libp2p/go-buffer-pool"
)

var log = logging.Logger("mplex")

var MaxMessageSize = 1 << 20

// Max time to block waiting for a slow reader to read from a stream before
// resetting it. Preferably, we'd have some form of back-pressure mechanism but
// we don't have that in this protocol.
var ReceiveTimeout = 5 * time.Second

// ErrShutdown is returned when operating on a shutdown session
var ErrShutdown = errors.New("session shut down")

// ErrTwoInitiators is returned when both sides think they're the initiator
var ErrTwoInitiators = errors.New("two initiators")

// ErrInvalidState is returned when the other side does something it shouldn't.
// In this case, we close the connection to be safe.
var ErrInvalidState = errors.New("received an unexpected message from the peer")

// +1 for initiator
const (
	newStreamTag = 0
	messageTag   = 2
	closeTag     = 4
	resetTag     = 6
)

type Multiplex struct {
	con       net.Conn
	buf       *bufio.Reader
	nextID    uint64
	initiator bool

	closed       chan struct{}
	shutdown     chan struct{}
	shutdownErr  error
	shutdownLock sync.Mutex

	wrTkn chan struct{}

	nstreams chan *Stream

	channels map[streamID]*Stream
	chLock   sync.Mutex
}

func NewMultiplex(con net.Conn, initiator bool) *Multiplex {
	mp := &Multiplex{
		con:       con,
		initiator: initiator,
		buf:       bufio.NewReader(con),
		channels:  make(map[streamID]*Stream),
		closed:    make(chan struct{}),
		shutdown:  make(chan struct{}),
		wrTkn:     make(chan struct{}, 1),
		nstreams:  make(chan *Stream, 16),
	}

	go mp.handleIncoming()

	mp.wrTkn <- struct{}{}

	return mp
}

func (mp *Multiplex) newStream(id streamID, name string) (s *Stream) {
	s = &Stream{
		id:          id,
		name:        name,
		dataIn:      make(chan []byte, 8),
		reset:       make(chan struct{}),
		mp:          mp,
	}

	s.closedLocal, s.doCloseLocal = context.WithCancel(context.Background())
	return
}

func (m *Multiplex) Accept() (*Stream, error) {
	select {
	case s, ok := <-m.nstreams:
		if !ok {
			return nil, errors.New("multiplex closed")
		}
		return s, nil
	case <-m.closed:
		return nil, m.shutdownErr
	}
}

func (mp *Multiplex) Close() error {
	mp.closeNoWait()

	// Wait for the receive loop to finish.
	<-mp.closed

	return nil
}

func (mp *Multiplex) closeNoWait() {
	mp.shutdownLock.Lock()
	select {
	case <-mp.shutdown:
	default:
		mp.con.Close()
		close(mp.shutdown)
	}
	mp.shutdownLock.Unlock()
}

func (mp *Multiplex) IsClosed() bool {
	select {
	case <-mp.closed:
		return true
	default:
		return false
	}
}

func (mp *Multiplex) sendMsg(ctx context.Context, header uint64, data []byte) error {
	select {
	case tkn := <-mp.wrTkn:
		defer func() { mp.wrTkn <- tkn }()
	case <-ctx.Done():
		return ctx.Err()
	}

	dl, hasDl := ctx.Deadline()
	if hasDl {
		if err := mp.con.SetWriteDeadline(dl); err != nil {
			return err
		}
	}
	buf := pool.Get(len(data) + 20)
	defer pool.Put(buf)

	n := 0
	n += binary.PutUvarint(buf[n:], header)
	n += binary.PutUvarint(buf[n:], uint64(len(data)))
	n += copy(buf[n:], data)

	written, err := mp.con.Write(buf[:n])
	if err != nil && written > 0 {
		// Bail. We've written partial message and can't do anything
		// about this.
		mp.con.Close()
		return err
	}

	if hasDl {
		// only return this error if we don't *already* have an error from the write.
		if err2 := mp.con.SetWriteDeadline(time.Time{}); err == nil && err2 != nil {
			return err2
		}
	}

	return err
}

func (mp *Multiplex) nextChanID() uint64 {
	out := mp.nextID
	mp.nextID++
	return out
}

func (mp *Multiplex) NewStream() (*Stream, error) {
	return mp.NewNamedStream("")
}

func (mp *Multiplex) NewNamedStream(name string) (*Stream, error) {
	mp.chLock.Lock()

	// We could call IsClosed but this is faster (given that we already have
	// the lock).
	if mp.channels == nil {
		return nil, ErrShutdown
	}

	sid := mp.nextChanID()
	header := (sid << 3) | newStreamTag

	if name == "" {
		name = fmt.Sprint(sid)
	}
	s := mp.newStream(streamID{
		id:        sid,
		initiator: true,
	}, name)
	mp.channels[s.id] = s
	mp.chLock.Unlock()

	err := mp.sendMsg(context.Background(), header, []byte(name))
	if err != nil {
		return nil, err
	}

	return s, nil
}

func (mp *Multiplex) cleanup() {
	mp.closeNoWait()
	mp.chLock.Lock()
	defer mp.chLock.Unlock()
	for _, msch := range mp.channels {
		msch.clLock.Lock()
		if !msch.closedRemote {
			msch.closedRemote = true
			// Cancel readers
			close(msch.reset)
		}

		msch.doCloseLocal();
		msch.clLock.Unlock()
	}
	// Don't remove this nil assignment. We check if this is nil to check if
	// the connection is closed when we already have the lock (faster than
	// checking if the stream is closed).
	mp.channels = nil
	if mp.shutdownErr == nil {
		mp.shutdownErr = ErrShutdown
	}
	close(mp.closed)
}

func (mp *Multiplex) handleIncoming() {
	defer mp.cleanup()

	recvTimeout := time.NewTimer(0)
	defer recvTimeout.Stop()

	if !recvTimeout.Stop() {
		<-recvTimeout.C
	}

	for {
		chID, tag, err := mp.readNextHeader()
		if err != nil {
			mp.shutdownErr = err
			return
		}

		remoteIsInitiator := tag&1 == 0
		ch := streamID{
			// true if *I'm* the initiator.
			initiator: !remoteIsInitiator,
			id:        chID,
		}
		// Rounds up the tag:
		// 0 -> 0
		// 1 -> 2
		// 2 -> 2
		// 3 -> 4
		// etc...
		tag += (tag & 1)

		b, err := mp.readNext()
		if err != nil {
			mp.shutdownErr = err
			return
		}

		mp.chLock.Lock()
		msch, ok := mp.channels[ch]
		mp.chLock.Unlock()

		switch tag {
		case newStreamTag:
			if ok {
				log.Debugf("received NewStream message for existing stream: %d", ch)
				mp.shutdownErr = ErrInvalidState
				return
			}

			name := string(b)
			pool.Put(b)

			msch = mp.newStream(ch, name)
			mp.chLock.Lock()
			mp.channels[ch] = msch
			mp.chLock.Unlock()
			select {
			case mp.nstreams <- msch:
			case <-mp.shutdown:
				return
			}

		case resetTag:
			if !ok {
				// This is *ok*. We forget the stream on reset.
				continue
			}
			msch.clLock.Lock()

			// Honestly, this check should never be true... It means we've leaked.
			// However, this is an error on *our* side so we shouldn't just bail.
			isClosed := msch.isClosed()
			if isClosed && msch.closedRemote {
				msch.clLock.Unlock()
				log.Errorf("leaked a completely closed stream")
				continue
			}

			if !msch.closedRemote {
				close(msch.reset)
			}
			msch.closedRemote = true
			msch.doCloseLocal()

			msch.clLock.Unlock()

			mp.chLock.Lock()
			delete(mp.channels, ch)
			mp.chLock.Unlock()
		case closeTag:
			if !ok {
				continue
			}

			msch.clLock.Lock()

			if msch.closedRemote {
				msch.clLock.Unlock()
				// Technically a bug on the other side. We
				// should consider killing the connection.
				continue
			}

			close(msch.dataIn)
			msch.closedRemote = true

			cleanup := msch.isClosed()

			msch.clLock.Unlock()

			if cleanup {
				mp.chLock.Lock()
				delete(mp.channels, ch)
				mp.chLock.Unlock()
			}
		case messageTag:
			if !ok {
				// reset stream, return b
				pool.Put(b)

				// This is a perfectly valid case when we reset
				// and forget about the stream.
				log.Debugf("message for non-existant stream, dropping data: %d", ch)
				go mp.sendMsg(context.Background(), ch.header(resetTag), nil)
				continue
			}

			msch.clLock.Lock()
			remoteClosed := msch.closedRemote
			msch.clLock.Unlock()
			if remoteClosed {
				// closed stream, return b
				pool.Put(b)

				log.Errorf("Received data from remote after stream was closed by them. (len = %d)", len(b))
				go mp.sendMsg(context.Background(), msch.id.header(resetTag), nil)
				continue
			}

			recvTimeout.Reset(ReceiveTimeout)
			select {
			case msch.dataIn <- b:
			case <-msch.reset:
				pool.Put(b)
			case <-recvTimeout.C:
				pool.Put(b)
				log.Warningf("timed out receiving message into stream queue.")
				// Do not do this asynchronously. Otherwise, we
				// could drop a message, then receive a message,
				// then reset.
				msch.Reset()
				continue
			case <-mp.shutdown:
				pool.Put(b)
				return
			}
			if !recvTimeout.Stop() {
				<-recvTimeout.C
			}
		default:
			log.Debugf("message with unknown header on stream %s", ch)
			if ok {
				msch.Reset()
			}
		}
	}
}

func (mp *Multiplex) readNextHeader() (uint64, uint64, error) {
	h, err := binary.ReadUvarint(mp.buf)
	if err != nil {
		return 0, 0, err
	}

	// get channel ID
	ch := h >> 3

	rem := h & 7

	return ch, rem, nil
}

func (mp *Multiplex) readNext() ([]byte, error) {
	// get length
	l, err := binary.ReadUvarint(mp.buf)
	if err != nil {
		return nil, err
	}

	if l > uint64(MaxMessageSize) {
		return nil, fmt.Errorf("message size too large!")
	}

	if l == 0 {
		return nil, nil
	}

	buf := pool.Get(int(l))
	n, err := io.ReadFull(mp.buf, buf)
	if err != nil {
		return nil, err
	}

	return buf[:n], nil
}
<commit_msg>release lock when failing to create a new stream<commit_after>package multiplex

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	pool "github.com/libp2p/go-buffer-pool"
)

var log = logging.Logger("mplex")

var MaxMessageSize = 1 << 20

// Max time to block waiting for a slow reader to read from a stream before
// resetting it. Preferably, we'd have some form of back-pressure mechanism but
// we don't have that in this protocol.
var ReceiveTimeout = 5 * time.Second

// ErrShutdown is returned when operating on a shutdown session
var ErrShutdown = errors.New("session shut down")

// ErrTwoInitiators is returned when both sides think they're the initiator
var ErrTwoInitiators = errors.New("two initiators")

// ErrInvalidState is returned when the other side does something it shouldn't.
// In this case, we close the connection to be safe.
var ErrInvalidState = errors.New("received an unexpected message from the peer")

// +1 for initiator
const (
	newStreamTag = 0
	messageTag   = 2
	closeTag     = 4
	resetTag     = 6
)

type Multiplex struct {
	con       net.Conn
	buf       *bufio.Reader
	nextID    uint64
	initiator bool

	closed       chan struct{}
	shutdown     chan struct{}
	shutdownErr  error
	shutdownLock sync.Mutex

	wrTkn chan struct{}

	nstreams chan *Stream

	channels map[streamID]*Stream
	chLock   sync.Mutex
}

func NewMultiplex(con net.Conn, initiator bool) *Multiplex {
	mp := &Multiplex{
		con:       con,
		initiator: initiator,
		buf:       bufio.NewReader(con),
		channels:  make(map[streamID]*Stream),
		closed:    make(chan struct{}),
		shutdown:  make(chan struct{}),
		wrTkn:     make(chan struct{}, 1),
		nstreams:  make(chan *Stream, 16),
	}

	go mp.handleIncoming()

	mp.wrTkn <- struct{}{}

	return mp
}

func (mp *Multiplex) newStream(id streamID, name string) (s *Stream) {
	s = &Stream{
		id:          id,
		name:        name,
		dataIn:      make(chan []byte, 8),
		reset:       make(chan struct{}),
		mp:          mp,
	}

	s.closedLocal, s.doCloseLocal = context.WithCancel(context.Background())
	return
}

func (m *Multiplex) Accept() (*Stream, error) {
	select {
	case s, ok := <-m.nstreams:
		if !ok {
			return nil, errors.New("multiplex closed")
		}
		return s, nil
	case <-m.closed:
		return nil, m.shutdownErr
	}
}

func (mp *Multiplex) Close() error {
	mp.closeNoWait()

	// Wait for the receive loop to finish.
	<-mp.closed

	return nil
}

func (mp *Multiplex) closeNoWait() {
	mp.shutdownLock.Lock()
	select {
	case <-mp.shutdown:
	default:
		mp.con.Close()
		close(mp.shutdown)
	}
	mp.shutdownLock.Unlock()
}

func (mp *Multiplex) IsClosed() bool {
	select {
	case <-mp.closed:
		return true
	default:
		return false
	}
}

func (mp *Multiplex) sendMsg(ctx context.Context, header uint64, data []byte) error {
	select {
	case tkn := <-mp.wrTkn:
		defer func() { mp.wrTkn <- tkn }()
	case <-ctx.Done():
		return ctx.Err()
	}

	dl, hasDl := ctx.Deadline()
	if hasDl {
		if err := mp.con.SetWriteDeadline(dl); err != nil {
			return err
		}
	}
	buf := pool.Get(len(data) + 20)
	defer pool.Put(buf)

	n := 0
	n += binary.PutUvarint(buf[n:], header)
	n += binary.PutUvarint(buf[n:], uint64(len(data)))
	n += copy(buf[n:], data)

	written, err := mp.con.Write(buf[:n])
	if err != nil && written > 0 {
		// Bail. We've written partial message and can't do anything
		// about this.
		mp.con.Close()
		return err
	}

	if hasDl {
		// only return this error if we don't *already* have an error from the write.
		if err2 := mp.con.SetWriteDeadline(time.Time{}); err == nil && err2 != nil {
			return err2
		}
	}

	return err
}

func (mp *Multiplex) nextChanID() uint64 {
	out := mp.nextID
	mp.nextID++
	return out
}

func (mp *Multiplex) NewStream() (*Stream, error) {
	return mp.NewNamedStream("")
}

func (mp *Multiplex) NewNamedStream(name string) (*Stream, error) {
	mp.chLock.Lock()

	// We could call IsClosed but this is faster (given that we already have
	// the lock).
	if mp.channels == nil {
		mp.chLock.Unlock()
		return nil, ErrShutdown
	}

	sid := mp.nextChanID()
	header := (sid << 3) | newStreamTag

	if name == "" {
		name = fmt.Sprint(sid)
	}
	s := mp.newStream(streamID{
		id:        sid,
		initiator: true,
	}, name)
	mp.channels[s.id] = s
	mp.chLock.Unlock()

	err := mp.sendMsg(context.Background(), header, []byte(name))
	if err != nil {
		return nil, err
	}

	return s, nil
}

func (mp *Multiplex) cleanup() {
	mp.closeNoWait()
	mp.chLock.Lock()
	defer mp.chLock.Unlock()
	for _, msch := range mp.channels {
		msch.clLock.Lock()
		if !msch.closedRemote {
			msch.closedRemote = true
			// Cancel readers
			close(msch.reset)
		}

		msch.doCloseLocal();
		msch.clLock.Unlock()
	}
	// Don't remove this nil assignment. We check if this is nil to check if
	// the connection is closed when we already have the lock (faster than
	// checking if the stream is closed).
	mp.channels = nil
	if mp.shutdownErr == nil {
		mp.shutdownErr = ErrShutdown
	}
	close(mp.closed)
}

func (mp *Multiplex) handleIncoming() {
	defer mp.cleanup()

	recvTimeout := time.NewTimer(0)
	defer recvTimeout.Stop()

	if !recvTimeout.Stop() {
		<-recvTimeout.C
	}

	for {
		chID, tag, err := mp.readNextHeader()
		if err != nil {
			mp.shutdownErr = err
			return
		}

		remoteIsInitiator := tag&1 == 0
		ch := streamID{
			// true if *I'm* the initiator.
			initiator: !remoteIsInitiator,
			id:        chID,
		}
		// Rounds up the tag:
		// 0 -> 0
		// 1 -> 2
		// 2 -> 2
		// 3 -> 4
		// etc...
		tag += (tag & 1)

		b, err := mp.readNext()
		if err != nil {
			mp.shutdownErr = err
			return
		}

		mp.chLock.Lock()
		msch, ok := mp.channels[ch]
		mp.chLock.Unlock()

		switch tag {
		case newStreamTag:
			if ok {
				log.Debugf("received NewStream message for existing stream: %d", ch)
				mp.shutdownErr = ErrInvalidState
				return
			}

			name := string(b)
			pool.Put(b)

			msch = mp.newStream(ch, name)
			mp.chLock.Lock()
			mp.channels[ch] = msch
			mp.chLock.Unlock()
			select {
			case mp.nstreams <- msch:
			case <-mp.shutdown:
				return
			}

		case resetTag:
			if !ok {
				// This is *ok*. We forget the stream on reset.
				continue
			}
			msch.clLock.Lock()

			// Honestly, this check should never be true... It means we've leaked.
			// However, this is an error on *our* side so we shouldn't just bail.
			isClosed := msch.isClosed()
			if isClosed && msch.closedRemote {
				msch.clLock.Unlock()
				log.Errorf("leaked a completely closed stream")
				continue
			}

			if !msch.closedRemote {
				close(msch.reset)
			}
			msch.closedRemote = true
			msch.doCloseLocal()

			msch.clLock.Unlock()

			mp.chLock.Lock()
			delete(mp.channels, ch)
			mp.chLock.Unlock()
		case closeTag:
			if !ok {
				continue
			}

			msch.clLock.Lock()

			if msch.closedRemote {
				msch.clLock.Unlock()
				// Technically a bug on the other side. We
				// should consider killing the connection.
				continue
			}

			close(msch.dataIn)
			msch.closedRemote = true

			cleanup := msch.isClosed()

			msch.clLock.Unlock()

			if cleanup {
				mp.chLock.Lock()
				delete(mp.channels, ch)
				mp.chLock.Unlock()
			}
		case messageTag:
			if !ok {
				// reset stream, return b
				pool.Put(b)

				// This is a perfectly valid case when we reset
				// and forget about the stream.
				log.Debugf("message for non-existant stream, dropping data: %d", ch)
				go mp.sendMsg(context.Background(), ch.header(resetTag), nil)
				continue
			}

			msch.clLock.Lock()
			remoteClosed := msch.closedRemote
			msch.clLock.Unlock()
			if remoteClosed {
				// closed stream, return b
				pool.Put(b)

				log.Errorf("Received data from remote after stream was closed by them. (len = %d)", len(b))
				go mp.sendMsg(context.Background(), msch.id.header(resetTag), nil)
				continue
			}

			recvTimeout.Reset(ReceiveTimeout)
			select {
			case msch.dataIn <- b:
			case <-msch.reset:
				pool.Put(b)
			case <-recvTimeout.C:
				pool.Put(b)
				log.Warningf("timed out receiving message into stream queue.")
				// Do not do this asynchronously. Otherwise, we
				// could drop a message, then receive a message,
				// then reset.
				msch.Reset()
				continue
			case <-mp.shutdown:
				pool.Put(b)
				return
			}
			if !recvTimeout.Stop() {
				<-recvTimeout.C
			}
		default:
			log.Debugf("message with unknown header on stream %s", ch)
			if ok {
				msch.Reset()
			}
		}
	}
}

func (mp *Multiplex) readNextHeader() (uint64, uint64, error) {
	h, err := binary.ReadUvarint(mp.buf)
	if err != nil {
		return 0, 0, err
	}

	// get channel ID
	ch := h >> 3

	rem := h & 7

	return ch, rem, nil
}

func (mp *Multiplex) readNext() ([]byte, error) {
	// get length
	l, err := binary.ReadUvarint(mp.buf)
	if err != nil {
		return nil, err
	}

	if l > uint64(MaxMessageSize) {
		return nil, fmt.Errorf("message size too large!")
	}

	if l == 0 {
		return nil, nil
	}

	buf := pool.Get(int(l))
	n, err := io.ReadFull(mp.buf, buf)
	if err != nil {
		return nil, err
	}

	return buf[:n], nil
}
<|endoftext|>


<commit_before>package nodeflock

import (
	"fmt"
	"os"
	"time"
)

type CallResult struct {
	Result interface{} `json:"result"`
	Error  string      `json:"error"`
	Log    []struct {
		Level string        `json:"level"`
		Data  []interface{} `json:"data"`
		Stack []string      `json:"stack"`
	} `json:"log"`
}
type processRentalApplicationResponse struct {
	err     error
	process *process
}

type processRentalApplication struct {
	chanResponse chan processRentalApplicationResponse
}

type Flock struct {
	size             int
	processes        []*process
	chanRentalReturn chan *process
	chanRentalApply  chan chan processRentalApplicationResponse //chan *processRentalApplication
	idCounter        int
}

func NewFlock(jsModuleFile string, size int, maxExecutionTime time.Duration) (f *Flock, err error) {
	f = &Flock{
		size:             size,
		processes:        []*process{},
		chanRentalReturn: make(chan *process),
		chanRentalApply:  make(chan chan processRentalApplicationResponse),
	}

	lastBusyReportTime := time.Now()
	reportAllBusy := func() {
		now := time.Now()
		if now.Sub(lastBusyReportTime) > 20*time.Second {
			lastBusyReportTime = now
			fmt.Println("everybody is busy")
		}
	}

	go func() {
		busy := map[int]bool{}
		applications := []chan processRentalApplicationResponse{} //[]*processRentalApplication{}
		chanProcessDied := make(chan int)
		lastSourceChange := getFileChange(jsModuleFile)
		for {
			// is anyone obsolote
			numOkProcesses := 0
			for _, deathRowProcess := range f.processes {
				if _, isBusy := busy[deathRowProcess.id]; !isBusy && deathRowProcess.sourceFileChange < lastSourceChange {
					deathRowProcess.kill()
				} else {
					numOkProcesses++
				}
			}

			if lastSourceChange > 0 {
				// spawn processes, if we lost some
			spawnLoop:
				for i := numOkProcesses; i < f.size; i++ {
					p, processErr := newProcess(jsModuleFile, lastSourceChange, maxExecutionTime, f.getId(), chanProcessDied)
					if processErr != nil {
						fmt.Println("breaking spawnloop could not launch procees", processErr)
						break spawnLoop
					}
					f.processes = append(f.processes, p)
				}
			}

			select {
			case <-time.After(time.Millisecond * 10):
				newSourceChange := getFileChange(jsModuleFile)
				if newSourceChange != lastSourceChange {
					fmt.Println("source file update", jsModuleFile, lastSourceChange, newSourceChange)
				}
				lastSourceChange = newSourceChange
				//fmt.Println("checking source", lastSourceChange)
			case deadProcessID := <-chanProcessDied:
				// a process died
				// handle error
				// remove process
				f.processes = filterProcesses(f.processes, func(process *process) bool {
					return process.id != deadProcessID
				})
				delete(busy, deadProcessID)
			case rentalApplication := <-f.chanRentalApply:
				//fmt.Println("incoming application")
				applications = append(applications, rentalApplication)
			case returnedProcess := <-f.chanRentalReturn:
				// back to the
				//fmt.Println("got one back")
				delete(busy, returnedProcess.id)
			}
			// let us see if anyone has time
			processedApplications := []chan processRentalApplicationResponse{}

		applicationLoop:
			for _, application := range applications {
				if len(busy) == len(f.processes) {
					// they are all busy
					reportAllBusy()
					break
				}
				var lonelyProcess *process
				lonelyProcess = nil
				for _, process := range f.processes {
					_, isBusy := busy[process.id]
					if !isBusy && process.sourceFileChange == lastSourceChange {
						if lonelyProcess == nil {
							lonelyProcess = process
						}
						if lonelyProcess.lastCall < process.lastCall {
							lonelyProcess = process
						}
					}
				}

				// this one is not busy
				processedApplications = append(processedApplications, application)
				// make the applicant happy
				if lonelyProcess != nil {
					busy[lonelyProcess.id] = true
					application <- processRentalApplicationResponse{
						process: lonelyProcess,
						err:     nil,
					}
					continue applicationLoop

				} else {
					fmt.Println("could not find lonely process")
				}
			}
			// clean up
			cleanedApplications := []chan processRentalApplicationResponse{}
		applicationLoopClean:
			for _, application := range applications {
				for _, processedApplication := range processedApplications {
					if processedApplication == application {
						continue applicationLoopClean
					}
				}
				cleanedApplications = append(cleanedApplications, application)
			}
			applications = cleanedApplications
			// fmt.Println(busy)
			//fmt.Println(time.Now().Sub(start))

		}
	}()

	return
}

func getFileChange(file string) int64 {
	fileInfo, err := os.Stat(file)
	if err != nil {
		return int64(0)
	}
	return fileInfo.ModTime().UnixNano()
}

func filterProcesses(processes []*process, filter func(*process) bool) (cleaned []*process) {
	survivors := []*process{}
	for _, process := range processes {
		if filter(process) {
			survivors = append(survivors, process)
		}
	}
	return survivors
}

func (f *Flock) getId() int {
	f.idCounter++
	return f.idCounter
}

func (f *Flock) borrowProcess() (p *process, err error) {
	// dummy implementation
	application := make(chan processRentalApplicationResponse)
	f.chanRentalApply <- application
	response := <-application
	return response.process, response.err
}

func (f *Flock) returnProcess(p *process) {
	//fmt.Println("returning process", p.id)
	f.chanRentalReturn <- p
}

func (f *Flock) CallRaw(jsonBytes []byte) (rawResult []byte, err error) {
	p, processErr := f.borrowProcess()
	if processErr != nil {
		err = processErr
		return
	}
	//fmt.Println("got a process", p.id)
	rawResult, callErr := p.rawCallJS(jsonBytes)
	f.returnProcess(p)
	return rawResult, callErr
}

func (f *Flock) Call(funcName string, args ...interface{}) (result CallResult, err error) {
	return
}
<commit_msg>fixed round robin, when determining a lonely process<commit_after>package nodeflock

import (
	"fmt"
	"os"
	"time"
)

type CallResult struct {
	Result interface{} `json:"result"`
	Error  string      `json:"error"`
	Log    []struct {
		Level string        `json:"level"`
		Data  []interface{} `json:"data"`
		Stack []string      `json:"stack"`
	} `json:"log"`
}
type processRentalApplicationResponse struct {
	err     error
	process *process
}

type processRentalApplication struct {
	chanResponse chan processRentalApplicationResponse
}

type Flock struct {
	size             int
	processes        []*process
	chanRentalReturn chan *process
	chanRentalApply  chan chan processRentalApplicationResponse //chan *processRentalApplication
	idCounter        int
}

func NewFlock(jsModuleFile string, size int, maxExecutionTime time.Duration) (f *Flock, err error) {
	f = &Flock{
		size:             size,
		processes:        []*process{},
		chanRentalReturn: make(chan *process),
		chanRentalApply:  make(chan chan processRentalApplicationResponse),
	}

	lastBusyReportTime := time.Now()
	reportAllBusy := func() {
		now := time.Now()
		if now.Sub(lastBusyReportTime) > 20*time.Second {
			lastBusyReportTime = now
			fmt.Println("everybody is busy")
		}
	}

	go func() {
		busy := map[int]bool{}
		applications := []chan processRentalApplicationResponse{} //[]*processRentalApplication{}
		chanProcessDied := make(chan int)
		lastSourceChange := getFileChange(jsModuleFile)
		for {
			// is anyone obsolete
			numOkProcesses := 0
			for _, deathRowProcess := range f.processes {
				if _, isBusy := busy[deathRowProcess.id]; !isBusy && deathRowProcess.sourceFileChange < lastSourceChange {
					deathRowProcess.kill()
				} else {
					numOkProcesses++
				}
			}

			if lastSourceChange > 0 {
				// spawn processes, if we lost some
			spawnLoop:
				for i := numOkProcesses; i < f.size; i++ {
					p, processErr := newProcess(jsModuleFile, lastSourceChange, maxExecutionTime, f.getId(), chanProcessDied)
					if processErr != nil {
						fmt.Println("breaking spawnloop could not launch procees", processErr)
						break spawnLoop
					}
					f.processes = append(f.processes, p)
				}
			}

			select {
			case <-time.After(time.Millisecond * 10):
				newSourceChange := getFileChange(jsModuleFile)
				if newSourceChange != lastSourceChange {
					fmt.Println("source file update", jsModuleFile, lastSourceChange, newSourceChange)
				}
				lastSourceChange = newSourceChange
				//fmt.Println("checking source", lastSourceChange)
			case deadProcessID := <-chanProcessDied:
				// a process died
				// handle error
				// remove process
				f.processes = filterProcesses(f.processes, func(process *process) bool {
					return process.id != deadProcessID
				})
				delete(busy, deadProcessID)
			case rentalApplication := <-f.chanRentalApply:
				//fmt.Println("incoming application")
				applications = append(applications, rentalApplication)
			case returnedProcess := <-f.chanRentalReturn:
				// back to the
				//fmt.Println("got one back")
				delete(busy, returnedProcess.id)
			}
			// let us see if anyone has time
			processedApplications := []chan processRentalApplicationResponse{}

		applicationLoop:
			for _, application := range applications {
				if len(busy) == len(f.processes) {
					// they are all busy
					reportAllBusy()
					break
				}
				var lonelyProcess *process
				lonelyProcess = nil
				//fmt.Println("looking for a process, that is not busy")
				for _, process := range f.processes {
					_, isBusy := busy[process.id]
					if !isBusy && process.sourceFileChange == lastSourceChange {
						if lonelyProcess == nil {
							lonelyProcess = process
						}
						//fmt.Println("lonely process:", lonelyProcess.id, lonelyProcess.lastCall, "process:", process.id, process.lastCall)
						if lonelyProcess.lastCall > process.lastCall {
							lonelyProcess = process
						}
					}
				}

				//fmt.Println("well that is my lonely process:", lonelyProcess.id, lonelyProcess.lastCall)
				// this one is not busy
				processedApplications = append(processedApplications, application)
				// make the applicant happy
				if lonelyProcess != nil {
					busy[lonelyProcess.id] = true
					application <- processRentalApplicationResponse{
						process: lonelyProcess,
						err:     nil,
					}
					continue applicationLoop

				} else {
					fmt.Println("could not find lonely process")
				}
			}
			// clean up
			cleanedApplications := []chan processRentalApplicationResponse{}
		applicationLoopClean:
			for _, application := range applications {
				for _, processedApplication := range processedApplications {
					if processedApplication == application {
						continue applicationLoopClean
					}
				}
				cleanedApplications = append(cleanedApplications, application)
			}
			applications = cleanedApplications
			// fmt.Println(busy)
			//fmt.Println(time.Now().Sub(start))

		}
	}()

	return
}

func getFileChange(file string) int64 {
	fileInfo, err := os.Stat(file)
	if err != nil {
		return int64(0)
	}
	return fileInfo.ModTime().UnixNano()
}

func filterProcesses(processes []*process, filter func(*process) bool) (cleaned []*process) {
	survivors := []*process{}
	for _, process := range processes {
		if filter(process) {
			survivors = append(survivors, process)
		}
	}
	return survivors
}

func (f *Flock) getId() int {
	f.idCounter++
	return f.idCounter
}

func (f *Flock) borrowProcess() (p *process, err error) {
	// dummy implementation
	application := make(chan processRentalApplicationResponse)
	f.chanRentalApply <- application
	response := <-application
	return response.process, response.err
}

func (f *Flock) returnProcess(p *process) {
	//fmt.Println("returning process", p.id)
	f.chanRentalReturn <- p
}

func (f *Flock) CallRaw(jsonBytes []byte) (rawResult []byte, err error) {
	p, processErr := f.borrowProcess()
	if processErr != nil {
		err = processErr
		return
	}
	//fmt.Println("got a process", p.id)
	rawResult, callErr := p.rawCallJS(jsonBytes)
	f.returnProcess(p)
	return rawResult, callErr
}

func (f *Flock) Call(funcName string, args ...interface{}) (result CallResult, err error) {
	return
}
<|endoftext|>


<commit_before>package publish

import (
	"strconv"

	schema "github.com/grafana/metrictank/schema"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	log "github.com/sirupsen/logrus"
)

var (
	ingestedMetrics = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gateway",
		Name:      "samples_ingested_total",
		Help:      "Number of samples ingested",
	}, []string{"org"})
)

type Publisher interface {
	Publish(metrics []*schema.MetricData) error
	Type() string
}

var (
	publisher Publisher
)

func Init(p Publisher) {
	if p == nil {
		publisher = &nullPublisher{}
	} else {
		publisher = p
	}
	log.Infof("using %s publisher", publisher.Type())
}

func Publish(metrics []*schema.MetricData) error {
	if len(metrics) == 0 {
		return nil
	}

	if err := publisher.Publish(metrics); err != nil {
		return err
	}
	// capture accounting data.
	orgCounts := make(map[int]int32)
	for _, m := range metrics {
		orgCounts[m.OrgId]++
	}
	for org, count := range orgCounts {
		ingestedMetrics.WithLabelValues(strconv.Itoa(org)).Add(float64(count))
	}
	return nil
}

// nullPublisher drops all metrics passed through the publish interface
type nullPublisher struct{}

func (*nullPublisher) Publish(metrics []*schema.MetricData) error {
	log.Debugf("publishing not enabled, dropping %d metrics", len(metrics))
	return nil
}

func (*nullPublisher) Type() string {
	return "nullPublisher"
}

func Persist(metrics []*schema.MetricData) error {
	return publisher.Publish(metrics)
}
<commit_msg>remove unused function<commit_after>package publish

import (
	"strconv"

	schema "github.com/grafana/metrictank/schema"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	log "github.com/sirupsen/logrus"
)

var (
	ingestedMetrics = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gateway",
		Name:      "samples_ingested_total",
		Help:      "Number of samples ingested",
	}, []string{"org"})
)

type Publisher interface {
	Publish(metrics []*schema.MetricData) error
	Type() string
}

var (
	publisher Publisher
)

func Init(p Publisher) {
	if p == nil {
		publisher = &nullPublisher{}
	} else {
		publisher = p
	}
	log.Infof("using %s publisher", publisher.Type())
}

func Publish(metrics []*schema.MetricData) error {
	if len(metrics) == 0 {
		return nil
	}

	if err := publisher.Publish(metrics); err != nil {
		return err
	}
	// capture accounting data.
	orgCounts := make(map[int]int32)
	for _, m := range metrics {
		orgCounts[m.OrgId]++
	}
	for org, count := range orgCounts {
		ingestedMetrics.WithLabelValues(strconv.Itoa(org)).Add(float64(count))
	}
	return nil
}

// nullPublisher drops all metrics passed through the publish interface
type nullPublisher struct{}

func (*nullPublisher) Publish(metrics []*schema.MetricData) error {
	log.Debugf("publishing not enabled, dropping %d metrics", len(metrics))
	return nil
}

func (*nullPublisher) Type() string {
	return "nullPublisher"
}
<|endoftext|>


<commit_before>package endpoint

import (
	"strconv"

	log "github.com/Sirupsen/logrus"
	"github.com/weaveworks/scope/probe/endpoint/procspy"
	"github.com/weaveworks/scope/probe/process"
	"github.com/weaveworks/scope/report"
)

// connectionTrackerConfig are the config options for the endpoint tracker.
type connectionTrackerConfig struct {
	HostID       string
	HostName     string
	SpyProcs     bool
	UseConntrack bool
	WalkProc     bool
	UseEbpfConn  bool
	ProcRoot     string
	BufferSize   int
	ProcessCache *process.CachingWalker
	Scanner      procspy.ConnectionScanner
	DNSSnooper   *DNSSnooper
}

type connectionTracker struct {
	conf            connectionTrackerConfig
	flowWalker      flowWalker // Interface
	ebpfTracker     eventTracker
	reverseResolver *reverseResolver
}

func newConnectionTracker(conf connectionTrackerConfig) connectionTracker {
	ct := connectionTracker{
		conf:            conf,
		reverseResolver: newReverseResolver(),
	}
	if conf.UseEbpfConn {
		et, err := newEbpfTracker()
		if err == nil {
			ct.ebpfTracker = et
			go ct.getInitialState()
			return ct
		}
		log.Warnf("Error setting up the eBPF tracker, falling back to proc scanning: %v", err)
	}
	ct.useProcfs()
	return ct
}

func flowToTuple(f flow) (ft fourTuple) {
	ft = fourTuple{
		f.Original.Layer3.SrcIP,
		f.Original.Layer3.DstIP,
		uint16(f.Original.Layer4.SrcPort),
		uint16(f.Original.Layer4.DstPort),
	}
	// Handle DNAT-ed connections in the initial state
	if f.Original.Layer3.DstIP != f.Reply.Layer3.SrcIP {
		ft = fourTuple{
			f.Reply.Layer3.DstIP,
			f.Reply.Layer3.SrcIP,
			uint16(f.Reply.Layer4.DstPort),
			uint16(f.Reply.Layer4.SrcPort),
		}
	}
	return ft
}

func (t *connectionTracker) useProcfs() {
	t.ebpfTracker = nil
	if t.conf.WalkProc && t.conf.Scanner == nil {
		t.conf.Scanner = procspy.NewConnectionScanner(t.conf.ProcessCache, t.conf.SpyProcs)
	}
	if t.flowWalker == nil {
		t.flowWalker = newConntrackFlowWalker(t.conf.UseConntrack, t.conf.ProcRoot, t.conf.BufferSize)
	}
}

// ReportConnections calls trackers according to the configuration.
func (t *connectionTracker) ReportConnections(rpt *report.Report) {
	hostNodeID := report.MakeHostNodeID(t.conf.HostID)

	if t.ebpfTracker != nil {
		if !t.ebpfTracker.isDead() {
			t.performEbpfTrack(rpt, hostNodeID)
			return
		}
		log.Warnf("ebpf tracker died, gently falling back to proc scanning")
		t.useProcfs()
	}

	// seenTuples contains information about connections seen by conntrack and it will be passed to the /proc parser
	seenTuples := map[string]fourTuple{}
	t.performFlowWalk(rpt, &seenTuples)
	if t.conf.WalkProc && t.conf.Scanner != nil {
		t.performWalkProc(rpt, hostNodeID, &seenTuples)
	}
}

func (t *connectionTracker) performFlowWalk(rpt *report.Report, seenTuples *map[string]fourTuple) {
	// Consult the flowWalker for short-lived connections
	extraNodeInfo := map[string]string{
		Conntracked: "true",
	}
	t.flowWalker.walkFlows(func(f flow, alive bool) {
		tuple := flowToTuple(f)
		(*seenTuples)[tuple.key()] = tuple
		t.addConnection(rpt, tuple, "", extraNodeInfo, extraNodeInfo)
	})
}

func (t *connectionTracker) performWalkProc(rpt *report.Report, hostNodeID string, seenTuples *map[string]fourTuple) error {
	conns, err := t.conf.Scanner.Connections()
	if err != nil {
		return err
	}
	for conn := conns.Next(); conn != nil; conn = conns.Next() {
		tuple, namespaceID, incoming := connectionTuple(conn, *seenTuples)
		var (
			toNodeInfo   = map[string]string{Procspied: "true"}
			fromNodeInfo = map[string]string{Procspied: "true"}
		)
		if conn.Proc.PID > 0 {
			fromNodeInfo[process.PID] = strconv.FormatUint(uint64(conn.Proc.PID), 10)
			fromNodeInfo[report.HostNodeID] = hostNodeID
		}
		if incoming {
			tuple.reverse()
			toNodeInfo, fromNodeInfo = fromNodeInfo, toNodeInfo
		}
		t.addConnection(rpt, tuple, namespaceID, fromNodeInfo, toNodeInfo)
	}
	return nil
}

// getInitialState runs conntrack and proc parsing synchronously only
// once to initialize ebpfTracker
func (t *connectionTracker) getInitialState() {
	var processCache *process.CachingWalker
	walker := process.NewWalker(t.conf.ProcRoot, true)
	processCache = process.NewCachingWalker(walker)
	processCache.Tick()

	scanner := procspy.NewSyncConnectionScanner(processCache, t.conf.SpyProcs)
	seenTuples := map[string]fourTuple{}
	// Consult the flowWalker to get the initial state
	if err := IsConntrackSupported(t.conf.ProcRoot); t.conf.UseConntrack && err != nil {
		log.Warnf("Not using conntrack: not supported by the kernel: %s", err)
	} else if existingFlows, err := existingConnections([]string{"--any-nat"}); err != nil {
		log.Errorf("conntrack existingConnections error: %v", err)
	} else {
		for _, f := range existingFlows {
			tuple := flowToTuple(f)
			seenTuples[tuple.key()] = tuple
		}
	}

	conns, err := scanner.Connections()
	if err != nil {
		log.Errorf("Error initializing ebpfTracker while scanning /proc, continuing without initial connections: %s", err)
	}
	scanner.Stop()

	processesWaitingInAccept := []int{}
	processCache.Walk(func(p, prev process.Process) {
		if p.IsWaitingInAccept {
			processesWaitingInAccept = append(processesWaitingInAccept, p.PID)
		}
	})

	t.ebpfTracker.feedInitialConnections(conns, seenTuples, processesWaitingInAccept, report.MakeHostNodeID(t.conf.HostID))
}

func (t *connectionTracker) performEbpfTrack(rpt *report.Report, hostNodeID string) error {
	t.ebpfTracker.walkConnections(func(e ebpfConnection) {
		fromNodeInfo := map[string]string{
			EBPF: "true",
		}
		toNodeInfo := map[string]string{
			EBPF: "true",
		}
		if e.pid > 0 {
			fromNodeInfo[process.PID] = strconv.Itoa(e.pid)
			fromNodeInfo[report.HostNodeID] = hostNodeID
		}

		if e.incoming {
			t.addConnection(rpt, reverse(e.tuple), e.networkNamespace, toNodeInfo, fromNodeInfo)
		} else {
			t.addConnection(rpt, e.tuple, e.networkNamespace, fromNodeInfo, toNodeInfo)
		}

	})
	return nil
}

func (t *connectionTracker) addConnection(rpt *report.Report, ft fourTuple, namespaceID string, extraFromNode, extraToNode map[string]string) {
	var (
		fromNode = t.makeEndpointNode(namespaceID, ft.fromAddr, ft.fromPort, extraFromNode)
		toNode   = t.makeEndpointNode(namespaceID, ft.toAddr, ft.toPort, extraToNode)
	)
	rpt.Endpoint = rpt.Endpoint.AddNode(fromNode.WithEdge(toNode.ID, report.EdgeMetadata{}))
	rpt.Endpoint = rpt.Endpoint.AddNode(toNode)
}

func (t *connectionTracker) makeEndpointNode(namespaceID string, addr string, port uint16, extra map[string]string) report.Node {
	portStr := strconv.Itoa(int(port))
	node := report.MakeNodeWith(report.MakeEndpointNodeID(t.conf.HostID, namespaceID, addr, portStr), nil)
	if names := t.conf.DNSSnooper.CachedNamesForIP(addr); len(names) > 0 {
		node = node.WithSet(SnoopedDNSNames, report.MakeStringSet(names...))
	}
	if names, err := t.reverseResolver.get(addr); err == nil && len(names) > 0 {
		node = node.WithSet(ReverseDNSNames, report.MakeStringSet(names...))
	}
	if extra != nil {
		node = node.WithLatests(extra)
	}
	return node
}

func (t *connectionTracker) Stop() error {
	if t.ebpfTracker != nil {
		t.ebpfTracker.stop()
	}
	if t.flowWalker != nil {
		t.flowWalker.stop()
	}
	t.reverseResolver.stop()
	return nil
}

func connectionTuple(conn *procspy.Connection, seenTuples map[string]fourTuple) (fourTuple, string, bool) {
	namespaceID := ""
	tuple := fourTuple{
		conn.LocalAddress.String(),
		conn.RemoteAddress.String(),
		conn.LocalPort,
		conn.RemotePort,
	}
	if conn.Proc.NetNamespaceID > 0 {
		namespaceID = strconv.FormatUint(conn.Proc.NetNamespaceID, 10)
	}

	// If we've already seen this connection, we should know the direction
	// (or have already figured it out), so we normalize and use the
	// canonical direction. Otherwise, we can use a port-heuristic to guess
	// the direction.
	canonical, ok := seenTuples[tuple.key()]
	return tuple, namespaceID, (ok && canonical != tuple) || (!ok && tuple.fromPort < tuple.toPort)
}
<commit_msg>refactor: remove superfluous pointering<commit_after>package endpoint

import (
	"strconv"

	log "github.com/Sirupsen/logrus"
	"github.com/weaveworks/scope/probe/endpoint/procspy"
	"github.com/weaveworks/scope/probe/process"
	"github.com/weaveworks/scope/report"
)

// connectionTrackerConfig are the config options for the endpoint tracker.
type connectionTrackerConfig struct {
	HostID       string
	HostName     string
	SpyProcs     bool
	UseConntrack bool
	WalkProc     bool
	UseEbpfConn  bool
	ProcRoot     string
	BufferSize   int
	ProcessCache *process.CachingWalker
	Scanner      procspy.ConnectionScanner
	DNSSnooper   *DNSSnooper
}

type connectionTracker struct {
	conf            connectionTrackerConfig
	flowWalker      flowWalker // Interface
	ebpfTracker     eventTracker
	reverseResolver *reverseResolver
}

func newConnectionTracker(conf connectionTrackerConfig) connectionTracker {
	ct := connectionTracker{
		conf:            conf,
		reverseResolver: newReverseResolver(),
	}
	if conf.UseEbpfConn {
		et, err := newEbpfTracker()
		if err == nil {
			ct.ebpfTracker = et
			go ct.getInitialState()
			return ct
		}
		log.Warnf("Error setting up the eBPF tracker, falling back to proc scanning: %v", err)
	}
	ct.useProcfs()
	return ct
}

func flowToTuple(f flow) (ft fourTuple) {
	ft = fourTuple{
		f.Original.Layer3.SrcIP,
		f.Original.Layer3.DstIP,
		uint16(f.Original.Layer4.SrcPort),
		uint16(f.Original.Layer4.DstPort),
	}
	// Handle DNAT-ed connections in the initial state
	if f.Original.Layer3.DstIP != f.Reply.Layer3.SrcIP {
		ft = fourTuple{
			f.Reply.Layer3.DstIP,
			f.Reply.Layer3.SrcIP,
			uint16(f.Reply.Layer4.DstPort),
			uint16(f.Reply.Layer4.SrcPort),
		}
	}
	return ft
}

func (t *connectionTracker) useProcfs() {
	t.ebpfTracker = nil
	if t.conf.WalkProc && t.conf.Scanner == nil {
		t.conf.Scanner = procspy.NewConnectionScanner(t.conf.ProcessCache, t.conf.SpyProcs)
	}
	if t.flowWalker == nil {
		t.flowWalker = newConntrackFlowWalker(t.conf.UseConntrack, t.conf.ProcRoot, t.conf.BufferSize)
	}
}

// ReportConnections calls trackers according to the configuration.
func (t *connectionTracker) ReportConnections(rpt *report.Report) {
	hostNodeID := report.MakeHostNodeID(t.conf.HostID)

	if t.ebpfTracker != nil {
		if !t.ebpfTracker.isDead() {
			t.performEbpfTrack(rpt, hostNodeID)
			return
		}
		log.Warnf("ebpf tracker died, gently falling back to proc scanning")
		t.useProcfs()
	}

	// seenTuples contains information about connections seen by conntrack and it will be passed to the /proc parser
	seenTuples := map[string]fourTuple{}
	t.performFlowWalk(rpt, seenTuples)
	if t.conf.WalkProc && t.conf.Scanner != nil {
		t.performWalkProc(rpt, hostNodeID, seenTuples)
	}
}

func (t *connectionTracker) performFlowWalk(rpt *report.Report, seenTuples map[string]fourTuple) {
	// Consult the flowWalker for short-lived connections
	extraNodeInfo := map[string]string{
		Conntracked: "true",
	}
	t.flowWalker.walkFlows(func(f flow, alive bool) {
		tuple := flowToTuple(f)
		seenTuples[tuple.key()] = tuple
		t.addConnection(rpt, tuple, "", extraNodeInfo, extraNodeInfo)
	})
}

func (t *connectionTracker) performWalkProc(rpt *report.Report, hostNodeID string, seenTuples map[string]fourTuple) error {
	conns, err := t.conf.Scanner.Connections()
	if err != nil {
		return err
	}
	for conn := conns.Next(); conn != nil; conn = conns.Next() {
		tuple, namespaceID, incoming := connectionTuple(conn, seenTuples)
		var (
			toNodeInfo   = map[string]string{Procspied: "true"}
			fromNodeInfo = map[string]string{Procspied: "true"}
		)
		if conn.Proc.PID > 0 {
			fromNodeInfo[process.PID] = strconv.FormatUint(uint64(conn.Proc.PID), 10)
			fromNodeInfo[report.HostNodeID] = hostNodeID
		}
		if incoming {
			tuple.reverse()
			toNodeInfo, fromNodeInfo = fromNodeInfo, toNodeInfo
		}
		t.addConnection(rpt, tuple, namespaceID, fromNodeInfo, toNodeInfo)
	}
	return nil
}

// getInitialState runs conntrack and proc parsing synchronously only
// once to initialize ebpfTracker
func (t *connectionTracker) getInitialState() {
	var processCache *process.CachingWalker
	walker := process.NewWalker(t.conf.ProcRoot, true)
	processCache = process.NewCachingWalker(walker)
	processCache.Tick()

	scanner := procspy.NewSyncConnectionScanner(processCache, t.conf.SpyProcs)
	seenTuples := map[string]fourTuple{}
	// Consult the flowWalker to get the initial state
	if err := IsConntrackSupported(t.conf.ProcRoot); t.conf.UseConntrack && err != nil {
		log.Warnf("Not using conntrack: not supported by the kernel: %s", err)
	} else if existingFlows, err := existingConnections([]string{"--any-nat"}); err != nil {
		log.Errorf("conntrack existingConnections error: %v", err)
	} else {
		for _, f := range existingFlows {
			tuple := flowToTuple(f)
			seenTuples[tuple.key()] = tuple
		}
	}

	conns, err := scanner.Connections()
	if err != nil {
		log.Errorf("Error initializing ebpfTracker while scanning /proc, continuing without initial connections: %s", err)
	}
	scanner.Stop()

	processesWaitingInAccept := []int{}
	processCache.Walk(func(p, prev process.Process) {
		if p.IsWaitingInAccept {
			processesWaitingInAccept = append(processesWaitingInAccept, p.PID)
		}
	})

	t.ebpfTracker.feedInitialConnections(conns, seenTuples, processesWaitingInAccept, report.MakeHostNodeID(t.conf.HostID))
}

func (t *connectionTracker) performEbpfTrack(rpt *report.Report, hostNodeID string) error {
	t.ebpfTracker.walkConnections(func(e ebpfConnection) {
		fromNodeInfo := map[string]string{
			EBPF: "true",
		}
		toNodeInfo := map[string]string{
			EBPF: "true",
		}
		if e.pid > 0 {
			fromNodeInfo[process.PID] = strconv.Itoa(e.pid)
			fromNodeInfo[report.HostNodeID] = hostNodeID
		}

		if e.incoming {
			t.addConnection(rpt, reverse(e.tuple), e.networkNamespace, toNodeInfo, fromNodeInfo)
		} else {
			t.addConnection(rpt, e.tuple, e.networkNamespace, fromNodeInfo, toNodeInfo)
		}

	})
	return nil
}

func (t *connectionTracker) addConnection(rpt *report.Report, ft fourTuple, namespaceID string, extraFromNode, extraToNode map[string]string) {
	var (
		fromNode = t.makeEndpointNode(namespaceID, ft.fromAddr, ft.fromPort, extraFromNode)
		toNode   = t.makeEndpointNode(namespaceID, ft.toAddr, ft.toPort, extraToNode)
	)
	rpt.Endpoint = rpt.Endpoint.AddNode(fromNode.WithEdge(toNode.ID, report.EdgeMetadata{}))
	rpt.Endpoint = rpt.Endpoint.AddNode(toNode)
}

func (t *connectionTracker) makeEndpointNode(namespaceID string, addr string, port uint16, extra map[string]string) report.Node {
	portStr := strconv.Itoa(int(port))
	node := report.MakeNodeWith(report.MakeEndpointNodeID(t.conf.HostID, namespaceID, addr, portStr), nil)
	if names := t.conf.DNSSnooper.CachedNamesForIP(addr); len(names) > 0 {
		node = node.WithSet(SnoopedDNSNames, report.MakeStringSet(names...))
	}
	if names, err := t.reverseResolver.get(addr); err == nil && len(names) > 0 {
		node = node.WithSet(ReverseDNSNames, report.MakeStringSet(names...))
	}
	if extra != nil {
		node = node.WithLatests(extra)
	}
	return node
}

func (t *connectionTracker) Stop() error {
	if t.ebpfTracker != nil {
		t.ebpfTracker.stop()
	}
	if t.flowWalker != nil {
		t.flowWalker.stop()
	}
	t.reverseResolver.stop()
	return nil
}

func connectionTuple(conn *procspy.Connection, seenTuples map[string]fourTuple) (fourTuple, string, bool) {
	namespaceID := ""
	tuple := fourTuple{
		conn.LocalAddress.String(),
		conn.RemoteAddress.String(),
		conn.LocalPort,
		conn.RemotePort,
	}
	if conn.Proc.NetNamespaceID > 0 {
		namespaceID = strconv.FormatUint(conn.Proc.NetNamespaceID, 10)
	}

	// If we've already seen this connection, we should know the direction
	// (or have already figured it out), so we normalize and use the
	// canonical direction. Otherwise, we can use a port-heuristic to guess
	// the direction.
	canonical, ok := seenTuples[tuple.key()]
	return tuple, namespaceID, (ok && canonical != tuple) || (!ok && tuple.fromPort < tuple.toPort)
}
<|endoftext|>


<commit_before>package api

import (
	"database/sql"
	"errors"
	"time"

	"github.com/pebblescape/pebblescape/Godeps/_workspace/src/github.com/jmoiron/sqlx"
	"github.com/pebblescape/pebblescape/pkg/utils"
)

var (
	ErrAppNameInvalid = errors.New("App name contains invalid characters. Alphanumeric only, no whitespace.")
	ErrAppNotExist    = errors.New("App does not exist")
	ErrAppNameTooLong = errors.New("App name too long. Max characters 30.")
	AppNameMax        = 30
)

// App represents the app in the database.
type App struct {
	ID        string     `db:"id"`
	Name      string     `db:"name"`
	CreatedAt *time.Time `db:"created_at"`
	UpdatedAt *time.Time `db:"updated_at"`
	DeletedAt *time.Time `db:"deleted_at"`
}

// AppRepo handles all interfacing between the host API and database.
type AppRepo struct {
	API *API
	DB  *sqlx.DB
}

// GetAppRepo creates and returns a new AppRepo instance.
func (a *API) GetAppRepo() *AppRepo {
	return &AppRepo{a, a.DB}
}

// List returns all apps in the database.
func (r *AppRepo) List() ([]App, error) {
	apps := []App{}

	if err := r.DB.Select(&apps, "SELECT * FROM apps ORDER BY name ASC"); err != nil {
		return apps, err
	}

	return apps, nil
}

// Get finds an app by ID and returns it.
func (r *AppRepo) Get(id string) (*App, error) {
	app := &App{}

	if err := r.DB.Get(app, "SELECT * FROM apps WHERE id=$1", id); err != nil {
		if err == sql.ErrNoRows {
			return app, ErrAppNotExist
		}
		return app, err
	}

	return app, nil
}

// GetByName finds an app by name and returns it.
func (r *AppRepo) GetByName(name string) (*App, error) {
	app := &App{}

	if err := r.DB.Get(app, "SELECT * FROM apps WHERE name=$1", name); err != nil {
		if err == sql.ErrNoRows {
			return app, ErrAppNotExist
		}
		return app, err
	}

	return app, nil
}

// Create saves a new app into the database.
func (r *AppRepo) Create(app *App) error {
	if !utils.AppNamePattern.Match([]byte(app.Name)) {
		return ErrAppNameInvalid
	}

	if len(app.Name) > AppNameMax {
		return ErrAppNameTooLong
	}

	if _, err := r.DB.NamedExec(`INSERT INTO apps (name) VALUES (:name)`, app); err != nil {
		return err
	}

	newApp, err := r.GetByName(app.Name)
	if err != nil {
		return err
	}

	app.ID = newApp.ID
	app.CreatedAt = newApp.CreatedAt
	app.UpdatedAt = newApp.UpdatedAt

	return nil
}

// Update changes values of an existing app.
func (r *AppRepo) Update(id string, data map[string]interface{}) (*App, error) {
	app, err := r.Get(id)
	if err != nil {
		return nil, err
	}

	return app, nil
}
<commit_msg>Comments<commit_after>package api

import (
	"database/sql"
	"errors"
	"time"

	"github.com/pebblescape/pebblescape/Godeps/_workspace/src/github.com/jmoiron/sqlx"
	"github.com/pebblescape/pebblescape/pkg/utils"
)

var (
	// ErrAppNameInvalid is the error returned when when trying to set invalid app name.
	ErrAppNameInvalid = errors.New("Name contains invalid characters. Alphanumeric only, no whitespace.")

	// ErrAppNotExist is the error returned when requested app ID or name doesn't exist.
	ErrAppNotExist = errors.New("App does not exist")

	// ErrAppNameTooLong is the error returned when trying to set too long app name.
	ErrAppNameTooLong = errors.New("App name too long. Max characters 30.")

	// AppNameMax is the maximum allowed length of an app name.
	AppNameMax = 30
)

// App represents the app in the database.
type App struct {
	ID        string     `db:"id"`
	Name      string     `db:"name"`
	CreatedAt *time.Time `db:"created_at"`
	UpdatedAt *time.Time `db:"updated_at"`
	DeletedAt *time.Time `db:"deleted_at"`
}

// AppRepo handles all interfacing between the host API and database.
type AppRepo struct {
	API *API
	DB  *sqlx.DB
}

// GetAppRepo creates and returns a new AppRepo instance.
func (a *API) GetAppRepo() *AppRepo {
	return &AppRepo{a, a.DB}
}

// List returns all apps in the database.
func (r *AppRepo) List() ([]App, error) {
	apps := []App{}

	if err := r.DB.Select(&apps, "SELECT * FROM apps ORDER BY name ASC"); err != nil {
		return apps, err
	}

	return apps, nil
}

// Get finds an app by ID and returns it.
func (r *AppRepo) Get(id string) (*App, error) {
	app := &App{}

	if err := r.DB.Get(app, "SELECT * FROM apps WHERE id=$1", id); err != nil {
		if err == sql.ErrNoRows {
			return app, ErrAppNotExist
		}
		return app, err
	}

	return app, nil
}

// GetByName finds an app by name and returns it.
func (r *AppRepo) GetByName(name string) (*App, error) {
	app := &App{}

	if err := r.DB.Get(app, "SELECT * FROM apps WHERE name=$1", name); err != nil {
		if err == sql.ErrNoRows {
			return app, ErrAppNotExist
		}
		return app, err
	}

	return app, nil
}

// Create saves a new app into the database.
func (r *AppRepo) Create(app *App) error {
	if !utils.AppNamePattern.Match([]byte(app.Name)) {
		return ErrAppNameInvalid
	}

	if len(app.Name) > AppNameMax {
		return ErrAppNameTooLong
	}

	if _, err := r.DB.NamedExec(`INSERT INTO apps (name) VALUES (:name)`, app); err != nil {
		return err
	}

	newApp, err := r.GetByName(app.Name)
	if err != nil {
		return err
	}

	app.ID = newApp.ID
	app.CreatedAt = newApp.CreatedAt
	app.UpdatedAt = newApp.UpdatedAt

	return nil
}

// Update changes values of an existing app.
func (r *AppRepo) Update(id string, data map[string]interface{}) (*App, error) {
	app, err := r.Get(id)
	if err != nil {
		return nil, err
	}

	return app, nil
}
<|endoftext|>


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

package service

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

type sequence struct {
	// The current position of the producer or consumer
	cursor,

	// The previous known position of the consumer (if producer) or producer (if consumer)
	gate,

	// These are fillers to pad the cache line, which is generally 64 bytes
	p2, p3, p4, p5, p6, p7 int64
}

func newSequence() *sequence {
	return &sequence{}
}

func (this *sequence) get() int64 {
	return atomic.LoadInt64(&this.cursor)
}

func (this *sequence) set(seq int64) {
	atomic.StoreInt64(&this.cursor, seq)
}

type buffer struct {
	id int64

	readIndex  int64 //读序号
	writeIndex int64 //写序号
	buf        []*[]byte

	size int64
	mask int64

	done int64

	pcond *sync.Cond
	ccond *sync.Cond
}

func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}

	if size == 0 {
		size = DefaultBufferSize
	}

	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		id:         atomic.AddInt64(&bufcnt, 1),
		readIndex:  int64(0),
		writeIndex: int64(0),
		buf:        make([]*[]byte, size),
		size:       size,
		mask:       size - 1,
		pcond:      sync.NewCond(new(sync.Mutex)),
		ccond:      sync.NewCond(new(sync.Mutex)),
	}, nil
}

/**
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

func (this *buffer) ID() int64 {
	return this.id
}

func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)

	this.pcond.L.Lock()
	this.ccond.Broadcast()
	this.pcond.L.Unlock()

	this.ccond.L.Lock()
	this.pcond.Broadcast()
	this.ccond.L.Unlock()

	return nil
}

/**
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() (p *[]byte, ok bool) {
	this.ccond.L.Lock()
	defer func() {
		this.pcond.Broadcast()
		this.ccond.L.Unlock()
		time.Sleep(5 * time.Microsecond)
	}()
	ok = false
	p = nil
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return nil, false
		}
		writeIndex = this.GetCurrentWriteIndex()
		if readIndex >= writeIndex {
			//fmt.Println("read wait")
			this.pcond.Broadcast()
			this.ccond.Wait()
		} else {
			break
		}
	}
	index := readIndex & this.mask //替代求模
	p = this.buf[index]
	this.buf[index] = nil
	atomic.AddInt64(&this.readIndex, int64(1))
	if p != nil {
		ok = true
	}
	return p, ok
}

/**
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in *[]byte) (ok bool) {
	this.pcond.L.Lock()
	defer func() {
		this.ccond.Broadcast()
		this.pcond.L.Unlock()
		time.Sleep(5 * time.Microsecond)
	}()
	ok = false
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return false
		}
		readIndex = this.GetCurrentReadIndex()
		if writeIndex >= readIndex && writeIndex-readIndex >= this.size {
			//fmt.Println("write wait")
			this.ccond.Broadcast()
			this.pcond.Wait()
			//time.Sleep(1 * time.Millisecond)
		} else {
			break
		}

	}
	index := writeIndex & this.mask //替代求模
	this.buf[index] = in
	atomic.AddInt64(&this.writeIndex, int64(1))
	ok = true
	return ok
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		time.Sleep(5 * time.Millisecond)
		if this.isDone() {
			return total, io.EOF
		}

		var b *[]byte
		default_size := int64(4096)

		b__ := make([]byte, default_size)
		b = &b__
		n, err := r.Read((*b)[0:1])
		if err != nil {
			return total, io.EOF
		}
		total += int64(n)
		max_cnt := 1
		for {
			if this.isDone() {
				return total, io.EOF
			}
			// If we have read 5 bytes and still not done, then there's a problem.
			if max_cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}
			_, err := r.Read((*b)[max_cnt:(max_cnt + 1)])

			//fmt.Println(b)
			if err != nil {
				return total, err
			}
			if (*b)[max_cnt] >= 0x80 {
				max_cnt++
			} else {
				break
			}
		}
		remlen, m := binary.Uvarint((*b)[1 : max_cnt+1])
		remlen_tmp := int64(remlen)
		start_ := int64(1) + int64(m)
		total_tmp := remlen_tmp + start_
		if total_tmp > default_size {
			write_bytes := make([]byte, 0, total_tmp)
			write_bytes = append(write_bytes, (*b)[0:m+1]...)
			b = &write_bytes
		}

		nlen := int64(start_)
		times := 0
		cnt_ := 32
		for nlen < total_tmp {
			if this.isDone() {
				return total, io.EOF
			}
			if times > 100 {
				return total, io.EOF
			} else {
				times = 0
			}
			times++
			tmpm := total_tmp - nlen

			var b_ []byte
			if tmpm < int64(cnt_) {
				b_ = (*b)[nlen:total_tmp]
			} else {
				b_ = (*b)[nlen : nlen+int64(cnt_)]
			}

			//b_ := make([]byte, remlen)
			n, err = r.Read(b_[0:])

			if err != nil {
				/*Log.Errorc(func() string {
					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
				})
				time.Sleep(5 * time.Millisecond)
				continue*/
				return total, err
			}
			//write_bytes = append(write_bytes, b_[0:]...)
			nlen += int64(n)
			total += int64(n)
		}

		ok := this.WriteBuffer(b)

		if !ok {
			return total, errors.New("write ringbuffer failed")
		}
	}
}

//func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
//	defer this.Close()
//
//	total := int64(0)
//
//	for {
//		time.Sleep(5 * time.Millisecond)
//		if this.isDone() {
//			return total, io.EOF
//		}
//
//		var write_bytes []byte
//
//		b := make([]byte, 5)
//		n, err := r.Read(b[0:1])
//		if err != nil {
//			return total, io.EOF
//		}
//		total += int64(n)
//		max_cnt := 1
//		for {
//			if this.isDone() {
//				return total, io.EOF
//			}
//			// If we have read 5 bytes and still not done, then there's a problem.
//			if max_cnt > 4 {
//				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
//			}
//			_, err := r.Read(b[max_cnt:(max_cnt + 1)])
//
//			//fmt.Println(b)
//			if err != nil {
//				return total, err
//			}
//			if b[max_cnt] >= 0x80 {
//				max_cnt++
//			} else {
//				break
//			}
//		}
//		remlen, m := binary.Uvarint(b[1 : max_cnt + 1])
//		remlen_tmp := int64(remlen)
//		total_tmp := remlen_tmp + int64(1) + int64(m)
//
//		write_bytes = make([]byte, 0, total_tmp)
//		write_bytes = append(write_bytes, b[0:m + 1]...)
//		nlen := int64(0)
//		times := 0
//		cnt_ := 32
//		for nlen < remlen_tmp {
//			if this.isDone() {
//				return total, io.EOF
//			}
//			if times > 100 {
//				return total, io.EOF
//			} else {
//				times = 0
//			}
//			times++
//			tmpm := remlen_tmp - nlen
//
//			var b_ []byte
//			if tmpm < int64(cnt_) {
//				b_ = make([]byte, tmpm)
//			} else {
//				b_ = make([]byte, cnt_)
//			}
//
//			//b_ := make([]byte, remlen)
//			n, err = r.Read(b_[0:])
//
//			if err != nil {
//				/*Log.Errorc(func() string {
//					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
//				})
//				time.Sleep(5 * time.Millisecond)
//				continue*/
//				return total, err
//			}
//			write_bytes = append(write_bytes, b_[0:]...)
//			nlen += int64(n)
//			total += int64(n)
//		}
//
//		ok := this.WriteBuffer(&write_bytes)
//
//		if !ok {
//			return total, errors.New("write ringbuffer failed")
//		}
//	}
//}

func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		if this.isDone() {
			return total, io.EOF
		}

		p, ok := this.ReadBuffer()
		if !ok {
			return total, errors.New("read buffer failed")
		}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			//Log.Debugc(func() string{ return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)})

			if err != nil {
				return total, err
			}
		}

	}
}

func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func ringCopy(dst, src []byte, start int64) int {
	n := len(src)

	i, l := 0, 0

	for n > 0 {
		l = copy(dst[start:], src[i:])
		i += l
		n -= l

		if n > 0 {
			start = 0
		}
	}

	return i
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<commit_msg>修改 buffer.go process.go sendrecv.go buffer modify sleep 5 Microsecond modify buffre.go readfrom sleep 5毫秒 xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

type sequence struct {
	// The current position of the producer or consumer
	cursor,

	// The previous known position of the consumer (if producer) or producer (if consumer)
	gate,

	// These are fillers to pad the cache line, which is generally 64 bytes
	p2, p3, p4, p5, p6, p7 int64
}

func newSequence() *sequence {
	return &sequence{}
}

func (this *sequence) get() int64 {
	return atomic.LoadInt64(&this.cursor)
}

func (this *sequence) set(seq int64) {
	atomic.StoreInt64(&this.cursor, seq)
}

type buffer struct {
	id int64

	readIndex  int64 //读序号
	writeIndex int64 //写序号
	buf        []*[]byte

	size int64
	mask int64

	done int64

	pcond *sync.Cond
	ccond *sync.Cond
}

func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}

	if size == 0 {
		size = DefaultBufferSize
	}

	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		id:         atomic.AddInt64(&bufcnt, 1),
		readIndex:  int64(0),
		writeIndex: int64(0),
		buf:        make([]*[]byte, size),
		size:       size,
		mask:       size - 1,
		pcond:      sync.NewCond(new(sync.Mutex)),
		ccond:      sync.NewCond(new(sync.Mutex)),
	}, nil
}

/**
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

func (this *buffer) ID() int64 {
	return this.id
}

func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)

	this.pcond.L.Lock()
	this.ccond.Broadcast()
	this.pcond.L.Unlock()

	this.ccond.L.Lock()
	this.pcond.Broadcast()
	this.ccond.L.Unlock()

	return nil
}

/**
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() (p *[]byte, ok bool) {
	this.ccond.L.Lock()
	defer func() {
		this.pcond.Broadcast()
		this.ccond.L.Unlock()
		time.Sleep(5 * time.Microsecond)
	}()
	ok = false
	p = nil
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return nil, false
		}
		writeIndex = this.GetCurrentWriteIndex()
		if readIndex >= writeIndex {
			//fmt.Println("read wait")
			this.pcond.Broadcast()
			this.ccond.Wait()
		} else {
			break
		}
	}
	index := readIndex & this.mask //替代求模
	p = this.buf[index]
	this.buf[index] = nil
	atomic.AddInt64(&this.readIndex, int64(1))
	if p != nil {
		ok = true
	}
	return p, ok
}

/**
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in *[]byte) (ok bool) {
	this.pcond.L.Lock()
	defer func() {
		this.ccond.Broadcast()
		this.pcond.L.Unlock()
		time.Sleep(5 * time.Microsecond)
	}()
	ok = false
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return false
		}
		readIndex = this.GetCurrentReadIndex()
		if writeIndex >= readIndex && writeIndex-readIndex >= this.size {
			//fmt.Println("write wait")
			this.ccond.Broadcast()
			this.pcond.Wait()
			//time.Sleep(1 * time.Millisecond)
		} else {
			break
		}

	}
	index := writeIndex & this.mask //替代求模
	this.buf[index] = in
	atomic.AddInt64(&this.writeIndex, int64(1))
	ok = true
	return ok
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		time.Sleep(5 * time.Millisecond)
		if this.isDone() {
			return total, io.EOF
		}

		var b *[]byte
		default_size := int64(4096)

		b__ := make([]byte, default_size)
		b = &b__
		n, err := r.Read((*b)[0:1])
		if err != nil {
			return total, io.EOF
		}
		total += int64(n)
		max_cnt := 1
		for {
			if this.isDone() {
				return total, io.EOF
			}
			// If we have read 5 bytes and still not done, then there's a problem.
			if max_cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}
			_, err := r.Read((*b)[max_cnt:(max_cnt + 1)])

			//fmt.Println(b)
			if err != nil {
				return total, err
			}
			if (*b)[max_cnt] >= 0x80 {
				max_cnt++
			} else {
				break
			}
		}
		remlen, m := binary.Uvarint((*b)[1 : max_cnt+1])
		remlen_tmp := int64(remlen)
		start_ := int64(1) + int64(m)
		total_tmp := remlen_tmp + start_
		if total_tmp > default_size {
			write_bytes := make([]byte, total_tmp)
			for i := int64(0); i < start_; i++ {
				write_bytes[i] = (*b)[i]
			}
			b = &write_bytes
		}

		nlen := int64(start_)
		times := 0
		cnt_ := 32
		for nlen < total_tmp {
			if this.isDone() {
				return total, io.EOF
			}
			if times > 100 {
				return total, io.EOF
			} else {
				times = 0
			}
			times++
			tmpm := total_tmp - nlen

			var b_ []byte
			if tmpm < int64(cnt_) {
				b_ = (*b)[nlen:total_tmp]
			} else {
				b_ = (*b)[nlen : nlen+int64(cnt_)]
			}

			//b_ := make([]byte, remlen)
			n, err = r.Read(b_[0:])

			if err != nil {
				/*Log.Errorc(func() string {
					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
				})
				time.Sleep(5 * time.Millisecond)
				continue*/
				return total, err
			}
			//write_bytes = append(write_bytes, b_[0:]...)
			nlen += int64(n)
			total += int64(n)
		}

		ok := this.WriteBuffer(b)

		if !ok {
			return total, errors.New("write ringbuffer failed")
		}
	}
}

//func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
//	defer this.Close()
//
//	total := int64(0)
//
//	for {
//		time.Sleep(5 * time.Millisecond)
//		if this.isDone() {
//			return total, io.EOF
//		}
//
//		var write_bytes []byte
//
//		b := make([]byte, 5)
//		n, err := r.Read(b[0:1])
//		if err != nil {
//			return total, io.EOF
//		}
//		total += int64(n)
//		max_cnt := 1
//		for {
//			if this.isDone() {
//				return total, io.EOF
//			}
//			// If we have read 5 bytes and still not done, then there's a problem.
//			if max_cnt > 4 {
//				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
//			}
//			_, err := r.Read(b[max_cnt:(max_cnt + 1)])
//
//			//fmt.Println(b)
//			if err != nil {
//				return total, err
//			}
//			if b[max_cnt] >= 0x80 {
//				max_cnt++
//			} else {
//				break
//			}
//		}
//		remlen, m := binary.Uvarint(b[1 : max_cnt + 1])
//		remlen_tmp := int64(remlen)
//		total_tmp := remlen_tmp + int64(1) + int64(m)
//
//		write_bytes = make([]byte, 0, total_tmp)
//		write_bytes = append(write_bytes, b[0:m + 1]...)
//		nlen := int64(0)
//		times := 0
//		cnt_ := 32
//		for nlen < remlen_tmp {
//			if this.isDone() {
//				return total, io.EOF
//			}
//			if times > 100 {
//				return total, io.EOF
//			} else {
//				times = 0
//			}
//			times++
//			tmpm := remlen_tmp - nlen
//
//			var b_ []byte
//			if tmpm < int64(cnt_) {
//				b_ = make([]byte, tmpm)
//			} else {
//				b_ = make([]byte, cnt_)
//			}
//
//			//b_ := make([]byte, remlen)
//			n, err = r.Read(b_[0:])
//
//			if err != nil {
//				/*Log.Errorc(func() string {
//					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
//				})
//				time.Sleep(5 * time.Millisecond)
//				continue*/
//				return total, err
//			}
//			write_bytes = append(write_bytes, b_[0:]...)
//			nlen += int64(n)
//			total += int64(n)
//		}
//
//		ok := this.WriteBuffer(&write_bytes)
//
//		if !ok {
//			return total, errors.New("write ringbuffer failed")
//		}
//	}
//}

func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		if this.isDone() {
			return total, io.EOF
		}

		p, ok := this.ReadBuffer()
		if !ok {
			return total, errors.New("read buffer failed")
		}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			//Log.Debugc(func() string{ return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)})

			if err != nil {
				return total, err
			}
		}

	}
}

func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func ringCopy(dst, src []byte, start int64) int {
	n := len(src)

	i, l := 0, 0

	for n > 0 {
		l = copy(dst[start:], src[i:])
		i += l
		n -= l

		if n > 0 {
			start = 0
		}
	}

	return i
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<|endoftext|>


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

package service

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

type sequence struct {
	// The current position of the producer or consumer
	cursor,

	// The previous known position of the consumer (if producer) or producer (if consumer)
	gate,

	// These are fillers to pad the cache line, which is generally 64 bytes
	p2, p3, p4, p5, p6, p7 int64
}

func newSequence() *sequence {
	return &sequence{}
}

func (this *sequence) get() int64 {
	return atomic.LoadInt64(&this.cursor)
}

func (this *sequence) set(seq int64) {
	atomic.StoreInt64(&this.cursor, seq)
}

type buffer struct {
	id int64

	readIndex  int64 //读序号
	writeIndex int64 //写序号
	buf        []*[]byte

	size int64
	mask int64

	done int64

	pcond *sync.Cond
	ccond *sync.Cond

	b []byte //readfrom中的临时数组，为反复使用不必创建新数组
}

func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}

	if size == 0 {
		size = DefaultBufferSize
	}

	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		id:         atomic.AddInt64(&bufcnt, 1),
		readIndex:  int64(0),
		writeIndex: int64(0),
		buf:        make([]*[]byte, size),
		size:       size,
		mask:       size - 1,
		pcond:      sync.NewCond(new(sync.Mutex)),
		ccond:      sync.NewCond(new(sync.Mutex)),
		b:          make([]byte, 5),
	}, nil
}

/**
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

func (this *buffer) ID() int64 {
	return this.id
}

func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)

	this.pcond.L.Lock()
	this.ccond.Signal()
	//this.ccond.Broadcast()
	this.pcond.L.Unlock()

	this.ccond.L.Lock()
	this.pcond.Signal()
	//this.pcond.Broadcast()
	this.ccond.L.Unlock()

	return nil
}

/**
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() (p *[]byte, ok bool) {
	this.ccond.L.Lock()
	defer func() {
		this.pcond.Signal()
		//this.pcond.Broadcast()
		this.ccond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()

	ok = false
	p = nil
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return nil, false
		}
		writeIndex = this.GetCurrentWriteIndex()
		if readIndex >= writeIndex {
			this.pcond.Signal()
			//this.pcond.Broadcast()
			this.ccond.Wait()
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := readIndex & this.mask //替代求模

	p = this.buf[index]
	this.buf[index] = nil

	atomic.AddInt64(&this.readIndex, int64(1))
	if p != nil {
		ok = true
	}
	return p, ok
}

/**
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in *[]byte) (ok bool) {
	this.pcond.L.Lock()
	defer func() {
		this.ccond.Signal()
		//this.ccond.Broadcast()
		this.pcond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()

	ok = false
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return false
		}
		readIndex = this.GetCurrentReadIndex()
		if writeIndex >= readIndex && writeIndex-readIndex >= this.size {
			this.ccond.Signal()
			//this.ccond.Broadcast()
			this.pcond.Wait()
			//time.Sleep(1 * time.Millisecond)
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := writeIndex & this.mask //替代求模

	this.buf[index] = in

	atomic.AddInt64(&this.writeIndex, int64(1))
	ok = true
	return ok
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		time.Sleep(5 * time.Millisecond)
		if this.isDone() {
			return total, io.EOF
		}

		var write_bytes []byte

		//b := make([]byte, 5)
		n, err := r.Read(this.b[0:1])
		if err != nil {
			return total, io.EOF
		}
		total += int64(n)
		max_cnt := 1
		for {
			if this.isDone() {
				return total, io.EOF
			}
			// If we have read 5 bytes and still not done, then there's a problem.
			if max_cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}
			_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

			//fmt.Println(b)
			if err != nil {
				return total, err
			}
			if this.b[max_cnt] >= 0x80 {
				max_cnt++
			} else {
				break
			}
		}
		remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
		remlen_tmp := int64(remlen)
		start_ := int64(1) + int64(m)
		total_tmp := remlen_tmp + start_

		write_bytes = make([]byte, total_tmp)
		copy(write_bytes[0:m+1], this.b[0:m+1])
		nlen := int64(0)
		times := 0
		for nlen < remlen_tmp {
			if this.isDone() {
				return total, io.EOF
			}
			if times > 100 {
				return total, io.EOF
			} else {
				times = 0
			}
			times++

			b_ := write_bytes[(start_ + nlen):remlen_tmp]

			//b_ := make([]byte, remlen)
			n, err = r.Read(b_[0:])

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
				})
				time.Sleep(5 * time.Millisecond)
				continue
				//return total, err
			}
			//write_bytes = append(write_bytes, b_[0:]...)
			nlen += int64(n)
			total += int64(n)
		}

		ok := this.WriteBuffer(&write_bytes)

		if !ok {
			return total, errors.New("write ringbuffer failed")
		}
	}
}

func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		if this.isDone() {
			return total, io.EOF
		}

		p, ok := this.ReadBuffer()
		if !ok {
			return total, errors.New("read buffer failed")
		}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			//Log.Debugc(func() string{ return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)})

			if err != nil {
				return total, err
			}
		}

	}
}

func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func ringCopy(dst, src []byte, start int64) int {
	n := len(src)

	i, l := 0, 0

	for n > 0 {
		l = copy(dst[start:], src[i:])
		i += l
		n -= l

		if n > 0 {
			start = 0
		}
	}

	return i
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom_not_receiver(r io.Reader) (*[]byte, error) {

	total := int64(0)

	if this.isDone() {
		return nil, io.EOF
	}

	var write_bytes []byte

	n, err := r.Read(this.b[0:1])
	if err != nil {
		return nil, io.EOF
	}
	total += int64(n)
	max_cnt := 1
	for {
		if this.isDone() {
			return nil, io.EOF
		}
		// If we have read 5 bytes and still not done, then there's a problem.
		if max_cnt > 4 {
			return nil, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
		}
		_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

		//fmt.Println(b)
		if err != nil {
			return nil, err
		}
		if this.b[max_cnt] >= 0x80 {
			max_cnt++
		} else {
			break
		}
	}
	remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
	remlen_tmp := int64(remlen)
	start_ := int64(1) + int64(m)
	total_tmp := remlen_tmp + start_

	write_bytes = make([]byte, total_tmp)
	copy(write_bytes[0:m+1], this.b[0:m+1])
	nlen := int64(0)
	times := 0
	cnt_ := int64(32)
	for nlen < remlen_tmp {
		if this.isDone() {
			return nil, io.EOF
		}
		if times > 100 {
			return nil, io.EOF
		} else {
			times = 0
		}
		times++
		tmpm := remlen_tmp - nlen

		if tmpm > cnt_ {
			n, err = r.Read(write_bytes[(start_ + nlen):(start_ + nlen + cnt_)])
		} else {
			n, err = r.Read(write_bytes[(start_ + nlen):])
		}

		if err != nil {
			return nil, err
		}
		nlen += int64(n)
		total += int64(n)
	}

	return &write_bytes, nil
}
<commit_msg>修改 buffer.go process.go sendrecv.go xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 还原receiver并修改部分代码<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

var (
	bufcnt            int64
	DefaultBufferSize int64

	DeviceInBufferSize  int64
	DeviceOutBufferSize int64

	MasterInBufferSize  int64
	MasterOutBufferSize int64
)

type sequence struct {
	// The current position of the producer or consumer
	cursor,

	// The previous known position of the consumer (if producer) or producer (if consumer)
	gate,

	// These are fillers to pad the cache line, which is generally 64 bytes
	p2, p3, p4, p5, p6, p7 int64
}

func newSequence() *sequence {
	return &sequence{}
}

func (this *sequence) get() int64 {
	return atomic.LoadInt64(&this.cursor)
}

func (this *sequence) set(seq int64) {
	atomic.StoreInt64(&this.cursor, seq)
}

type buffer struct {
	id int64

	readIndex  int64 //读序号
	writeIndex int64 //写序号
	buf        []*[]byte

	size int64
	mask int64

	done int64

	pcond *sync.Cond
	ccond *sync.Cond

	b []byte //readfrom中的临时数组，为反复使用不必创建新数组
}

func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}

	if size == 0 {
		size = DefaultBufferSize
	}

	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		id:         atomic.AddInt64(&bufcnt, 1),
		readIndex:  int64(0),
		writeIndex: int64(0),
		buf:        make([]*[]byte, size),
		size:       size,
		mask:       size - 1,
		pcond:      sync.NewCond(new(sync.Mutex)),
		ccond:      sync.NewCond(new(sync.Mutex)),
		b:          make([]byte, 5),
	}, nil
}

/**
获取当前读序号
*/
func (this *buffer) GetCurrentReadIndex() int64 {
	return atomic.LoadInt64(&this.readIndex)
}

/**
获取当前写序号
*/
func (this *buffer) GetCurrentWriteIndex() int64 {
	return atomic.LoadInt64(&this.writeIndex)
}

func (this *buffer) ID() int64 {
	return this.id
}

func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)

	this.pcond.L.Lock()
	this.ccond.Signal()
	//this.ccond.Broadcast()
	this.pcond.L.Unlock()

	this.ccond.L.Lock()
	this.pcond.Signal()
	//this.pcond.Broadcast()
	this.ccond.L.Unlock()

	return nil
}

/**
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
*/
func (this *buffer) ReadBuffer() (p *[]byte, ok bool) {
	this.ccond.L.Lock()
	defer func() {
		this.pcond.Signal()
		//this.pcond.Broadcast()
		this.ccond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()

	ok = false
	p = nil
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return nil, false
		}
		writeIndex = this.GetCurrentWriteIndex()
		if readIndex >= writeIndex {
			this.pcond.Signal()
			//this.pcond.Broadcast()
			this.ccond.Wait()
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := readIndex & this.mask //替代求模

	p = this.buf[index]
	this.buf[index] = nil

	atomic.AddInt64(&this.readIndex, int64(1))
	if p != nil {
		ok = true
	}
	return p, ok
}

/**
写入ringbuffer指针，以及将写序号加1
*/
func (this *buffer) WriteBuffer(in *[]byte) (ok bool) {
	this.pcond.L.Lock()
	defer func() {
		this.ccond.Signal()
		//this.ccond.Broadcast()
		this.pcond.L.Unlock()
		//time.Sleep(3 * time.Millisecond)
	}()

	ok = false
	readIndex := this.GetCurrentReadIndex()
	writeIndex := this.GetCurrentWriteIndex()
	for {
		if this.isDone() {
			return false
		}
		readIndex = this.GetCurrentReadIndex()
		if writeIndex >= readIndex && writeIndex-readIndex >= this.size {
			this.ccond.Signal()
			//this.ccond.Broadcast()
			this.pcond.Wait()
			//time.Sleep(1 * time.Millisecond)
			//runtime.Gosched()
			//time.Sleep(5 * time.Millisecond)
		} else {
			break
		}
		//time.Sleep(500 * time.Microsecond)
		//time.Sleep(1 * time.Millisecond)
	}
	index := writeIndex & this.mask //替代求模

	this.buf[index] = in

	atomic.AddInt64(&this.writeIndex, int64(1))
	ok = true
	return ok
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		time.Sleep(5 * time.Millisecond)
		if this.isDone() {
			return total, io.EOF
		}

		var write_bytes []byte

		//b := make([]byte, 5)
		n, err := r.Read(this.b[0:1])
		if err != nil {
			return total, io.EOF
		}
		total += int64(n)
		max_cnt := 1
		for {
			if this.isDone() {
				return total, io.EOF
			}
			// If we have read 5 bytes and still not done, then there's a problem.
			if max_cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}
			_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

			//fmt.Println(b)
			if err != nil {
				return total, err
			}
			if this.b[max_cnt] >= 0x80 {
				max_cnt++
			} else {
				break
			}
		}
		remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
		remlen_tmp := int64(remlen)
		start_ := int64(1) + int64(m)
		total_tmp := remlen_tmp + start_

		write_bytes = make([]byte, total_tmp)
		copy(write_bytes[0:m+1], this.b[0:m+1])
		nlen := int64(0)
		times := 0
		for nlen < remlen_tmp {
			if this.isDone() {
				return total, io.EOF
			}
			if times > 100 {
				return total, io.EOF
			} else {
				times = 0
			}
			times++

			b_ := write_bytes[(start_ + nlen):remlen_tmp]

			//b_ := make([]byte, remlen)
			n, err = r.Read(b_[0:])

			if err != nil {
				/*Log.Errorc(func() string {
					return fmt.Sprintf("从conn读取数据失败(%s)(0)", err)
				})*/
				time.Sleep(5 * time.Millisecond)
				continue
				//return total, err
			}
			//write_bytes = append(write_bytes, b_[0:]...)
			nlen += int64(n)
			total += int64(n)
		}

		ok := this.WriteBuffer(&write_bytes)

		if !ok {
			return total, errors.New("write ringbuffer failed")
		}
	}
}

func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()

	total := int64(0)

	for {
		if this.isDone() {
			return total, io.EOF
		}

		p, ok := this.ReadBuffer()
		if !ok {
			return total, errors.New("read buffer failed")
		}
		// There's some data, let's process it first
		if len(*p) > 0 {
			n, err := w.Write(*p)
			total += int64(n)
			//Log.Debugc(func() string{ return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)})

			if err != nil {
				return total, err
			}
		}

	}
}

func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func ringCopy(dst, src []byte, start int64) int {
	n := len(src)

	i, l := 0, 0

	for n > 0 {
		l = copy(dst[start:], src[i:])
		i += l
		n -= l

		if n > 0 {
			start = 0
		}
	}

	return i
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n&(n-1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}

/**
修改 尽量减少数据的创建
*/
func (this *buffer) ReadFrom_not_receiver(r io.Reader) (*[]byte, error) {

	total := int64(0)

	if this.isDone() {
		return nil, io.EOF
	}

	var write_bytes []byte

	n, err := r.Read(this.b[0:1])
	if err != nil {
		return nil, io.EOF
	}
	total += int64(n)
	max_cnt := 1
	for {
		if this.isDone() {
			return nil, io.EOF
		}
		// If we have read 5 bytes and still not done, then there's a problem.
		if max_cnt > 4 {
			return nil, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
		}
		_, err := r.Read(this.b[max_cnt:(max_cnt + 1)])

		//fmt.Println(b)
		if err != nil {
			return nil, err
		}
		if this.b[max_cnt] >= 0x80 {
			max_cnt++
		} else {
			break
		}
	}
	remlen, m := binary.Uvarint(this.b[1 : max_cnt+1])
	remlen_tmp := int64(remlen)
	start_ := int64(1) + int64(m)
	total_tmp := remlen_tmp + start_

	write_bytes = make([]byte, total_tmp)
	copy(write_bytes[0:m+1], this.b[0:m+1])
	nlen := int64(0)
	times := 0
	cnt_ := int64(32)
	for nlen < remlen_tmp {
		if this.isDone() {
			return nil, io.EOF
		}
		if times > 100 {
			return nil, io.EOF
		} else {
			times = 0
		}
		times++
		tmpm := remlen_tmp - nlen

		if tmpm > cnt_ {
			n, err = r.Read(write_bytes[(start_ + nlen):(start_ + nlen + cnt_)])
		} else {
			n, err = r.Read(write_bytes[(start_ + nlen):])
		}

		if err != nil {
			return nil, err
		}
		nlen += int64(n)
		total += int64(n)
	}

	return &write_bytes, nil
}
<|endoftext|>