<commit_before>package internal

import (
	"../agent"
	"google.golang.org/grpc"
	"golang.org/x/net/context"
	"sync"
	"bytes"
	"io"
	"time"
	"errors"
	"net"
	"unsafe"
	"sync/atomic"
)

const (
	PipeChannelBuffSize int = 10
	PipeAcksMaxSize int = 100
)

func intMin(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

type agentStream interface {
	Send(*agent.DataPacket) error
	Recv() (*agent.DataPacket, error)
}

type unAck struct {
	no uint32
	t  time.Time
}

func newUnAck(no uint32) *unAck {
	return &unAck{
		no: no,
		t: time.Now(),
	}
}

const (
	cmdPushAck = iota
	cmdPullAck
	cmdPopUnack
	cmdPushUnack
)

type command struct {
	cmd     int
	inAcks  []uint32
	outAcks chan uint32
}

//wrap grpc stream as net.Conn
type StreamPipe struct {
	ctx         context.Context
	cancelFunc  context.CancelFunc

	raw         agentStream
	cc          *ClientConn //may is nil

	waitGroup   sync.WaitGroup
	ioComplete	chan int

	reads       chan []byte
	writes      chan []byte
	writeable	chan int
	writeFlush	chan int

	cmds		chan *command

	acks        []uint32
	unAcks      []*unAck

	wBuffer     bytes.Buffer
	rBuffer     bytes.Buffer

	err         unsafe.Pointer

	acksChecker *time.Ticker

	serial      uint32
}

func NewStreamPipe(stream agentStream) *StreamPipe {
	ctx, cancelFunc := context.WithCancel(context.Background())
	pipe := &StreamPipe{
		ctx: ctx,
		cancelFunc:cancelFunc,
		raw: stream,
		ioComplete: make(chan int, 2),
		reads: make(chan []byte, PipeChannelBuffSize),
		writes: make(chan []byte, PipeChannelBuffSize),
		writeable: make(chan int, PipeChannelBuffSize),
		writeFlush: make(chan int, 1),
		cmds: make(chan *command, PipeChannelBuffSize),
		acksChecker: time.NewTicker(defaultAckCheckDelay),
	}

	pipe.waitGroup.Add(3)
	pipe.ioComplete <- 1
	pipe.ioComplete <- 1
	go pipe.readLoop()
	go pipe.writeLoop()
	go pipe.loop()

	return pipe
}

func (pipe *StreamPipe) Attach(cc *ClientConn) {
	pipe.cc = cc
}

func (pipe *StreamPipe) incrSerial() uint32 {
	pipe.serial++
	pipe.serial = pipe.serial & ^uint32(0)
	return pipe.serial
}

func (pipe *StreamPipe) newPacket() *agent.DataPacket {
	return &agent.DataPacket{
		No: pipe.incrSerial(),
	}
}

func (pipe *StreamPipe) readLoop() {
	defer func() {
		pipe.waitGroup.Done()
		close(pipe.reads)
		<- pipe.ioComplete
	}()

	for {
		//util error(contain eof)
		packet, err := pipe.raw.Recv()
		if err != nil {
			pipe.cancel(err)
			break
		}

		//pop unack
		if len(packet.Acks) > 0 {
			cmd := &command{
				cmd: cmdPopUnack,
				inAcks: packet.Acks,
			}
			pipe.cmds <- cmd
		}

		//push ack and buff
		if len(packet.Buff) > 0 {
			cmd := &command{
				cmd: cmdPushAck,
				inAcks: []uint32{packet.No},
			}
			pipe.cmds <- cmd

			pipe.reads <- packet.Buff
		}
	}
}

//allow first == nil
func (pipe *StreamPipe) handleWrite(first []byte) (err error) {
	if first != nil {
		_, err = pipe.wBuffer.Write(first)
		if err != nil {
			return err
		}
	}

	//until no need write
	FIRST:
	for {
		packet := pipe.newPacket()
		cap := defaultPacketMaxBytes

		//until packet is full or no more data
		SECOND:
		for cap > 0 {
			if pipe.wBuffer.Len() > 0 {
				lenght := intMin(cap, pipe.wBuffer.Len())
				buff := make([]byte, lenght)

				var nr int
				nr, err = pipe.wBuffer.Read(buff)
				if err != nil {
					return err
				}
				packet.Buff = append(packet.Buff, buff[:nr]...)
				cap -= nr

				continue
			}

			//non-block
			select {
			case buff, ok := <- pipe.writes:
				if !ok {
					//pipe.writes is closed
					break SECOND
				}
				<- pipe.writeable
				_, err = pipe.wBuffer.Write(buff)
				if err != nil {
					return err
				}
			default:
				break SECOND
			}
		}

		//request ack
		cmd := &command{
			cmd: cmdPullAck,
			outAcks: make(chan uint32),
		}
		pipe.cmds <- cmd

		for ack := range cmd.outAcks {
			packet.Acks = append(packet.Acks, ack)
		}

		if len(packet.Buff) == 0 && len(packet.Acks) == 0 {
			//no need write
			break FIRST
		}

		if len(packet.Buff) > 0 {
			//write out
			cmd = &command{
				cmd: cmdPushUnack,
				inAcks: []uint32{packet.No},
			}
			pipe.cmds <- cmd
		}

		err = pipe.raw.Send(packet)
		if err != nil {
			return err
		}
	}

	return nil
}

func (pipe *StreamPipe) writeLoop() {
	defer func() {
		pipe.waitGroup.Done()
		<- pipe.ioComplete
	}()

	var err error = nil

	//util pipe.writes is closed
	FIRST:
	for {
		select {
		case buff, ok := <-pipe.writes:
			if !ok {
				//pipe.writes is closed
				break FIRST
			}

			<- pipe.writeable
			err = pipe.handleWrite(buff)
			if err != nil {
				pipe.cancel(err)
				break FIRST
			}
		case <- pipe.writeFlush:
			err = pipe.handleWrite(nil)
			if err != nil {
				pipe.cancel(err)
				break FIRST
			}
		}
	}

	if err != nil {
		//discard packet
		for _ = range pipe.writes {}
	} else {
		//client actively close the stream
		//server wait for the peer to close the stream
		if s, ok := pipe.raw.(grpc.ClientStream); ok {
			if err := s.CloseSend(); err != nil {
				pipe.cancel(err)
			}
		}
	}
}

func (pipe *StreamPipe) ackCheck() error {
	if len(pipe.unAcks) == 0 {
		return nil
	}

	t := pipe.unAcks[0].t
	if time.Now().Sub(t) > defaultAckMaxDelay {
		return ErrAckTimeout
	}

	return nil
}

func (pipe *StreamPipe) handleCommand(cmd *command) (err error) {
	switch cmd.cmd {
	case cmdPushAck:
		pipe.acks = append(pipe.acks, cmd.inAcks...)
		if len(pipe.writes) == 0 {
			//non-blocking
			go func() {
				select {
				case pipe.writeFlush <- 1:
				case <- pipe.ctx.Done():
				}
			}()
		}
	case cmdPullAck:<commit_msg>internal/StreamPipe: Do not repeat flush<commit_after>package internal

import (
	"../agent"
	"google.golang.org/grpc"
	"golang.org/x/net/context"
	"sync"
	"bytes"
	"io"
	"time"
	"errors"
	"net"
	"unsafe"
	"sync/atomic"
)

const (
	PipeChannelBuffSize int = 10
	PipeAcksMaxSize int = 100
)

func intMin(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

type agentStream interface {
	Send(*agent.DataPacket) error
	Recv() (*agent.DataPacket, error)
}

type unAck struct {
	no uint32
	t  time.Time
}

func newUnAck(no uint32) *unAck {
	return &unAck{
		no: no,
		t: time.Now(),
	}
}

const (
	cmdPushAck = iota
	cmdPullAck
	cmdPopUnack
	cmdPushUnack
)

type command struct {
	cmd     int
	inAcks  []uint32
	outAcks chan uint32
}

//wrap grpc stream as net.Conn
type StreamPipe struct {
	ctx         context.Context
	cancelFunc  context.CancelFunc

	raw         agentStream
	cc          *ClientConn //may is nil

	waitGroup   sync.WaitGroup
	ioComplete	chan int

	reads       chan []byte
	writes      chan []byte
	writeable	chan int
	writeFlush	chan int

	cmds		chan *command

	acks        []uint32
	unAcks      []*unAck

	wBuffer     bytes.Buffer
	rBuffer     bytes.Buffer

	err         unsafe.Pointer

	acksChecker *time.Ticker

	serial      uint32
}

func NewStreamPipe(stream agentStream) *StreamPipe {
	ctx, cancelFunc := context.WithCancel(context.Background())
	pipe := &StreamPipe{
		ctx: ctx,
		cancelFunc:cancelFunc,
		raw: stream,
		ioComplete: make(chan int, 2),
		reads: make(chan []byte, PipeChannelBuffSize),
		writes: make(chan []byte, PipeChannelBuffSize),
		writeable: make(chan int, PipeChannelBuffSize),
		writeFlush: make(chan int, 1),
		cmds: make(chan *command, PipeChannelBuffSize),
		acksChecker: time.NewTicker(defaultAckCheckDelay),
	}

	pipe.waitGroup.Add(3)
	pipe.ioComplete <- 1
	pipe.ioComplete <- 1
	go pipe.readLoop()
	go pipe.writeLoop()
	go pipe.loop()

	return pipe
}

func (pipe *StreamPipe) Attach(cc *ClientConn) {
	pipe.cc = cc
}

func (pipe *StreamPipe) incrSerial() uint32 {
	pipe.serial++
	pipe.serial = pipe.serial & ^uint32(0)
	return pipe.serial
}

func (pipe *StreamPipe) newPacket() *agent.DataPacket {
	return &agent.DataPacket{
		No: pipe.incrSerial(),
	}
}

func (pipe *StreamPipe) readLoop() {
	defer func() {
		pipe.waitGroup.Done()
		close(pipe.reads)
		<- pipe.ioComplete
	}()

	for {
		//util error(contain eof)
		packet, err := pipe.raw.Recv()
		if err != nil {
			pipe.cancel(err)
			break
		}

		//pop unack
		if len(packet.Acks) > 0 {
			cmd := &command{
				cmd: cmdPopUnack,
				inAcks: packet.Acks,
			}
			pipe.cmds <- cmd
		}

		//push ack and buff
		if len(packet.Buff) > 0 {
			cmd := &command{
				cmd: cmdPushAck,
				inAcks: []uint32{packet.No},
			}
			pipe.cmds <- cmd

			pipe.reads <- packet.Buff
		}
	}
}

//allow first == nil
func (pipe *StreamPipe) handleWrite(first []byte) (err error) {
	if first != nil {
		_, err = pipe.wBuffer.Write(first)
		if err != nil {
			return err
		}
	}

	//until no need write
	FIRST:
	for {
		packet := pipe.newPacket()
		cap := defaultPacketMaxBytes

		//until packet is full or no more data
		SECOND:
		for cap > 0 {
			if pipe.wBuffer.Len() > 0 {
				lenght := intMin(cap, pipe.wBuffer.Len())
				buff := make([]byte, lenght)

				var nr int
				nr, err = pipe.wBuffer.Read(buff)
				if err != nil {
					return err
				}
				packet.Buff = append(packet.Buff, buff[:nr]...)
				cap -= nr

				continue
			}

			//non-block
			select {
			case buff, ok := <- pipe.writes:
				if !ok {
					//pipe.writes is closed
					break SECOND
				}
				<- pipe.writeable
				_, err = pipe.wBuffer.Write(buff)
				if err != nil {
					return err
				}
			default:
				break SECOND
			}
		}

		//request ack
		cmd := &command{
			cmd: cmdPullAck,
			outAcks: make(chan uint32),
		}
		pipe.cmds <- cmd

		for ack := range cmd.outAcks {
			packet.Acks = append(packet.Acks, ack)
		}

		if len(packet.Buff) == 0 && len(packet.Acks) == 0 {
			//no need write
			break FIRST
		}

		if len(packet.Buff) > 0 {
			//write out
			cmd = &command{
				cmd: cmdPushUnack,
				inAcks: []uint32{packet.No},
			}
			pipe.cmds <- cmd
		}

		err = pipe.raw.Send(packet)
		if err != nil {
			return err
		}
	}

	return nil
}

func (pipe *StreamPipe) writeLoop() {
	defer func() {
		pipe.waitGroup.Done()
		<- pipe.ioComplete
	}()

	var err error = nil

	//util pipe.writes is closed
	FIRST:
	for {
		select {
		case buff, ok := <-pipe.writes:
			if !ok {
				//pipe.writes is closed
				break FIRST
			}

			<- pipe.writeable
			err = pipe.handleWrite(buff)
			if err != nil {
				pipe.cancel(err)
				break FIRST
			}
		case <- pipe.writeFlush:
			err = pipe.handleWrite(nil)
			if err != nil {
				pipe.cancel(err)
				break FIRST
			}
		}
	}

	if err != nil {
		//discard packet
		for _ = range pipe.writes {}
	} else {
		//client actively close the stream
		//server wait for the peer to close the stream
		if s, ok := pipe.raw.(grpc.ClientStream); ok {
			if err := s.CloseSend(); err != nil {
				pipe.cancel(err)
			}
		}
	}
}

func (pipe *StreamPipe) ackCheck() error {
	if len(pipe.unAcks) == 0 {
		return nil
	}

	t := pipe.unAcks[0].t
	if time.Now().Sub(t) > defaultAckMaxDelay {
		return ErrAckTimeout
	}

	return nil
}

func (pipe *StreamPipe) handleCommand(cmd *command) (err error) {
	switch cmd.cmd {
	case cmdPushAck:
		pipe.acks = append(pipe.acks, cmd.inAcks...)
		if len(pipe.writes) == 0 {
			//Do not repeat flush
			select {
			case pipe.writeFlush <- 1:
			default:
			}
		}
	case cmdPullAck:
{'fraction_non_alphanumeric': 0.13196081988028296, 'fraction_numerical': 0.007981135497914021, 'mean_word_length': 2.6213464696223316, 'pattern_counts': {'":': 0, '<': 36, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6839185', 'n_tokens_mistral': 4903, 'n_tokens_neox': 4184, 'n_words': 1065}