<commit_before>package pulsar

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	"github.com/pkg/errors"

	command "github.com/t2y/go-pulsar/proto/command"
)

const (
	writeChanSize = 32
	readChanSize  = 32
)

type Response struct {
	Frame *command.Frame
	Error error
}

type AsyncTcpConn struct {
	wch  chan proto.Message
	rch  chan *Response
	conn *net.TCPConn

	readMutex        sync.Mutex
	reqMutex         sync.Mutex
	sendReceiveMutex sync.Mutex
}

func (ac *AsyncTcpConn) write(data []byte) (total int, err error) {
	if _, err = io.Copy(ac.conn, bytes.NewBuffer(data)); err != nil {
		err = errors.Wrap(err, "failed to write to connection")
		return
	}
	return
}

func (ac *AsyncTcpConn) writeLoop() {
	for {
		msg, ok := <-ac.wch
		if !ok {
			return
		}

		data, err := command.NewMarshaledBase(msg)
		if err != nil {
			log.WithFields(log.Fields{
				"err": err,
			}).Error("failed to marshal message")
			continue
		}

		_, err = ac.write(data)
		if err != nil {
			log.WithFields(log.Fields{
				"err": err,
			}).Error("failed to write in writeLoop")
			continue
		}
	}
}

func (ac *AsyncTcpConn) readFrame(size int64) (frame *bytes.Buffer, err error) {
	frame = bytes.NewBuffer(make([]byte, 0, size))
	if _, err = io.CopyN(frame, ac.conn, size); err != nil {
		err = errors.Wrap(err, "failed to read frame")
		return
	}
	return
}

func (ac *AsyncTcpConn) read() (frame *command.Frame, err error) {
	// there are 2 framing formats.
	//
	// 1. simple: [TOTAL_SIZE] [CMD_SIZE] [CMD]
	//
	// 2. with payload: [TOTAL_SIZE] [CMD_SIZE][CMD]
	//					[MAGIC_NUMBER][CHECKSUM] [METADATA_SIZE][METADATA]
	//					[PAYLOAD]

	ac.readMutex.Lock()
	defer ac.readMutex.Unlock()

	totalSizeFrame, err := ac.readFrame(int64(command.FrameSizeFieldSize))
	if err != nil {
		err = errors.Wrap(err, "failed to read total size frame")
		return
	}

	totalSize := binary.BigEndian.Uint32(totalSizeFrame.Bytes())
	log.Debug(totalSize)

	cmdSizeFrame, err := ac.readFrame(int64(command.FrameSizeFieldSize))
	if err != nil {
		err = errors.Wrap(err, "failed to read command size frame")
		return
	}

	cmdSize := binary.BigEndian.Uint32(cmdSizeFrame.Bytes())
	log.Debug(cmdSize)

	cmdFrame, err := ac.readFrame(int64(cmdSize))
	if err != nil {
		err = errors.Wrap(err, "failed to read command body frame")
		return
	}

	frame = new(command.Frame)
	frame.Cmddata = cmdFrame.Bytes()

	otherFramesSize := totalSize - (cmdSize + command.FrameSizeFieldSize)
	if otherFramesSize > 0 {
		var _otherFrames *bytes.Buffer
		_otherFrames, err = ac.readFrame(int64(otherFramesSize))
		if err != nil {
			err = errors.Wrap(err, "failed to read other frames")
			return
		}
		otherFrames := _otherFrames.Bytes()

		magicNumber := otherFrames[0:command.FrameMagicNumberFieldSize]
		log.Debug(magicNumber)

		checksumPos := command.FrameMagicNumberFieldSize + command.FrameChecksumSize
		frame.Checksum = otherFrames[command.FrameMagicNumberFieldSize:checksumPos]
		log.Debug(frame.Checksum)

		metadataSizePos := checksumPos + command.FrameMetadataFieldSize
		metadataSize := binary.BigEndian.Uint32(otherFrames[checksumPos:metadataSizePos])
		log.Debug(metadataSize)

		metadataPos := metadataSizePos + metadataSize
		frame.Metadata = otherFrames[metadataSizePos:metadataPos]
		log.Debug(frame.Metadata)
		frame.Payload = otherFrames[metadataPos:]
		log.Debug(frame.Payload)
	}

	return
}

func (ac *AsyncTcpConn) readLoop() {
	for {
		frame, err := ac.read()
		if err != nil {
			switch e := errors.Cause(err); e {
			case io.EOF:
				return // maybe connection was closed
			default:
				err = errors.Wrap(err, "failed to read in readLoop")
			}
		}

		if ac.rch == nil {
			return
		}
		ac.rch <- &Response{Frame: frame, Error: err}
	}
}

func (ac *AsyncTcpConn) Send(msg proto.Message) {
	ac.sendReceiveMutex.Lock()
	defer ac.sendReceiveMutex.Unlock()

	ac.wch <- msg
}

func (ac *AsyncTcpConn) Receive() (frame *command.Frame, err error) {
	ac.sendReceiveMutex.Lock()
	defer ac.sendReceiveMutex.Unlock()

	response, ok := <-ac.rch
	if !ok {
		err = errors.New("read channel has closed")
		return
	}<commit_msg>refactoring mutex handling to optimize<commit_after>package pulsar

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	"github.com/pkg/errors"

	command "github.com/t2y/go-pulsar/proto/command"
)

const (
	writeChanSize = 32
	readChanSize  = 32
)

type Response struct {
	Frame *command.Frame
	Error error
}

type AsyncTcpConn struct {
	wch  chan proto.Message
	rch  chan *Response
	conn *net.TCPConn

	readMutex        sync.Mutex
	reqMutex         sync.Mutex
	sendReceiveMutex sync.Mutex
}

func (ac *AsyncTcpConn) write(data []byte) (total int, err error) {
	if _, err = io.Copy(ac.conn, bytes.NewBuffer(data)); err != nil {
		err = errors.Wrap(err, "failed to write to connection")
		return
	}
	return
}

func (ac *AsyncTcpConn) writeLoop() {
	for {
		msg, ok := <-ac.wch
		if !ok {
			return
		}

		data, err := command.NewMarshaledBase(msg)
		if err != nil {
			log.WithFields(log.Fields{
				"err": err,
			}).Error("failed to marshal message")
			continue
		}

		_, err = ac.write(data)
		if err != nil {
			log.WithFields(log.Fields{
				"err": err,
			}).Error("failed to write in writeLoop")
			continue
		}
	}
}

func (ac *AsyncTcpConn) readFrame(size int64) (frame *bytes.Buffer, err error) {
	frame = bytes.NewBuffer(make([]byte, 0, size))
	if _, err = io.CopyN(frame, ac.conn, size); err != nil {
		err = errors.Wrap(err, "failed to read frame")
		return
	}
	return
}

func (ac *AsyncTcpConn) read() (frame *command.Frame, err error) {
	// there are 2 framing formats.
	//
	// 1. simple: [TOTAL_SIZE] [CMD_SIZE] [CMD]
	//
	// 2. with payload: [TOTAL_SIZE] [CMD_SIZE][CMD]
	//					[MAGIC_NUMBER][CHECKSUM] [METADATA_SIZE][METADATA]
	//					[PAYLOAD]

	ac.readMutex.Lock()
	defer ac.readMutex.Unlock()

	totalSizeFrame, err := ac.readFrame(int64(command.FrameSizeFieldSize))
	if err != nil {
		err = errors.Wrap(err, "failed to read total size frame")
		return
	}

	totalSize := binary.BigEndian.Uint32(totalSizeFrame.Bytes())
	log.Debug(totalSize)

	cmdSizeFrame, err := ac.readFrame(int64(command.FrameSizeFieldSize))
	if err != nil {
		err = errors.Wrap(err, "failed to read command size frame")
		return
	}

	cmdSize := binary.BigEndian.Uint32(cmdSizeFrame.Bytes())
	log.Debug(cmdSize)

	cmdFrame, err := ac.readFrame(int64(cmdSize))
	if err != nil {
		err = errors.Wrap(err, "failed to read command body frame")
		return
	}

	frame = new(command.Frame)
	frame.Cmddata = cmdFrame.Bytes()

	otherFramesSize := totalSize - (cmdSize + command.FrameSizeFieldSize)
	if otherFramesSize > 0 {
		var _otherFrames *bytes.Buffer
		_otherFrames, err = ac.readFrame(int64(otherFramesSize))
		if err != nil {
			err = errors.Wrap(err, "failed to read other frames")
			return
		}
		otherFrames := _otherFrames.Bytes()

		magicNumber := otherFrames[0:command.FrameMagicNumberFieldSize]
		log.Debug(magicNumber)

		checksumPos := command.FrameMagicNumberFieldSize + command.FrameChecksumSize
		frame.Checksum = otherFrames[command.FrameMagicNumberFieldSize:checksumPos]
		log.Debug(frame.Checksum)

		metadataSizePos := checksumPos + command.FrameMetadataFieldSize
		metadataSize := binary.BigEndian.Uint32(otherFrames[checksumPos:metadataSizePos])
		log.Debug(metadataSize)

		metadataPos := metadataSizePos + metadataSize
		frame.Metadata = otherFrames[metadataSizePos:metadataPos]
		log.Debug(frame.Metadata)
		frame.Payload = otherFrames[metadataPos:]
		log.Debug(frame.Payload)
	}

	return
}

func (ac *AsyncTcpConn) readLoop() {
	for {
		frame, err := ac.read()
		if err != nil {
			switch e := errors.Cause(err); e {
			case io.EOF:
				return // maybe connection was closed
			default:
				err = errors.Wrap(err, "failed to read in readLoop")
			}
		}

		if ac.rch == nil {
			return
		}
		ac.rch <- &Response{Frame: frame, Error: err}
	}
}

func (ac *AsyncTcpConn) Send(msg proto.Message) {
	ac.sendReceiveMutex.Lock()
	ac.wch <- msg
	ac.sendReceiveMutex.Unlock()
}

func (ac *AsyncTcpConn) Receive() (frame *command.Frame, err error) {
	ac.sendReceiveMutex.Lock()
	response, ok := <-ac.rch
	ac.sendReceiveMutex.Unlock()

	if !ok {
		err = errors.New("read channel has closed")
		return
	}
{'fraction_non_alphanumeric': 0.13302972195589646, 'fraction_numerical': 0.006471716203259827, 'mean_word_length': 3.8404872389791183, 'pattern_counts': {'":': 4, '<': 11, '<?xml version=': 0, '>': 5, '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': '1298693', 'n_tokens_mistral': 3288, 'n_tokens_neox': 2978, 'n_words': 682}