<commit_before>package smux

import (
	"bytes"
	"sync"
	"sync/atomic"

	"github.com/pkg/errors"
)

// Stream implements io.ReadWriteCloser
type Stream struct {
	id          uint32
	rstflag     int32
	sess        *Session
	buffer      bytes.Buffer
	bufferLock  sync.Mutex
	frameSize   int
	chReadEvent chan struct{} // notify a read event
	die         chan struct{} // flag the stream has closed
	dieLock     sync.Mutex
}

// newStream initiates a Stream struct
func newStream(id uint32, frameSize int, sess *Session) *Stream {
	s := new(Stream)
	s.id = id
	s.chReadEvent = make(chan struct{}, 1)
	s.frameSize = frameSize
	s.sess = sess
	s.die = make(chan struct{})
	return s
}

// Read implements io.ReadWriteCloser
func (s *Stream) Read(b []byte) (n int, err error) {
READ:
	select {
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	default:
	}

	s.bufferLock.Lock()
	n, err = s.buffer.Read(b)
	s.bufferLock.Unlock()

	if n > 0 {
		s.sess.returnTokens(n)
		return n, nil
	} else if atomic.LoadInt32(&s.rstflag) == 1 {
		s.Close()
		return 0, errors.New(errBrokenPipe)
	}

	select {
	case <-s.chReadEvent:
		goto READ
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	}
}

// Write implements io.ReadWriteCloser
func (s *Stream) Write(b []byte) (n int, err error) {
	select {
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	default:
	}

	frames := s.split(b, cmdPSH, s.id)
	// combine the frames
	var buffer bytes.Buffer
	for k := range frames {
		bts, _ := frames[k].MarshalBinary()
		buffer.Write(bts)
	}

	if _, err = s.sess.writeBinary(buffer.Bytes()); err != nil {
		return 0, err
	}
	return len(b), nil
}

// Close implements io.ReadWriteCloser
func (s *Stream) Close() error {
	s.dieLock.Lock()
	defer s.dieLock.Unlock()

	select {
	case <-s.die:
		return errors.New(errBrokenPipe)
	default:
		close(s.die)
		s.sess.streamClosed(s.id)
		s.sess.writeFrame(newFrame(cmdRST, s.id))
	}
	return nil
}

// notify the stream that the session has closed
func (s *Stream) sessionClose() error {
	s.dieLock.Lock()
	defer s.dieLock.Unlock()

	select {
	case <-s.die:
		return errors.New(errBrokenPipe)
	default:
		close(s.die)
	}
	return nil
}

// pushBytes a slice into buffer
func (s *Stream) pushBytes(p []byte) {
	s.bufferLock.Lock()
	s.buffer.Write(p)
	s.bufferLock.Unlock()
}

// recycleTokens transform remaining bytes to tokens(will truncate buffer)
func (s *Stream) recycleTokens() (n int) {
	s.bufferLock.Lock()
	n = s.buffer.Len()
	s.buffer.Reset()
	s.bufferLock.Unlock()
	return
}

// split large byte buffer into smaller frames, reference only
func (s *Stream) split(bts []byte, cmd byte, sid uint32) []Frame {
	var frames []Frame
	for len(bts) > int(s.frameSize) {<commit_msg>optimize sessionClose<commit_after>package smux

import (
	"bytes"
	"sync"
	"sync/atomic"

	"github.com/pkg/errors"
)

// Stream implements io.ReadWriteCloser
type Stream struct {
	id          uint32
	rstflag     int32
	sess        *Session
	buffer      bytes.Buffer
	bufferLock  sync.Mutex
	frameSize   int
	chReadEvent chan struct{} // notify a read event
	die         chan struct{} // flag the stream has closed
	dieLock     sync.Mutex
}

// newStream initiates a Stream struct
func newStream(id uint32, frameSize int, sess *Session) *Stream {
	s := new(Stream)
	s.id = id
	s.chReadEvent = make(chan struct{}, 1)
	s.frameSize = frameSize
	s.sess = sess
	s.die = make(chan struct{})
	return s
}

// Read implements io.ReadWriteCloser
func (s *Stream) Read(b []byte) (n int, err error) {
READ:
	select {
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	default:
	}

	s.bufferLock.Lock()
	n, err = s.buffer.Read(b)
	s.bufferLock.Unlock()

	if n > 0 {
		s.sess.returnTokens(n)
		return n, nil
	} else if atomic.LoadInt32(&s.rstflag) == 1 {
		s.Close()
		return 0, errors.New(errBrokenPipe)
	}

	select {
	case <-s.chReadEvent:
		goto READ
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	}
}

// Write implements io.ReadWriteCloser
func (s *Stream) Write(b []byte) (n int, err error) {
	select {
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	default:
	}

	frames := s.split(b, cmdPSH, s.id)
	// combine the frames
	var buffer bytes.Buffer
	for k := range frames {
		bts, _ := frames[k].MarshalBinary()
		buffer.Write(bts)
	}

	if _, err = s.sess.writeBinary(buffer.Bytes()); err != nil {
		return 0, err
	}
	return len(b), nil
}

// Close implements io.ReadWriteCloser
func (s *Stream) Close() error {
	s.dieLock.Lock()
	defer s.dieLock.Unlock()

	select {
	case <-s.die:
		return errors.New(errBrokenPipe)
	default:
		close(s.die)
		s.sess.streamClosed(s.id)
		s.sess.writeFrame(newFrame(cmdRST, s.id))
	}
	return nil
}

// session closes the stream
func (s *Stream) sessionClose() {
	s.dieLock.Lock()
	defer s.dieLock.Unlock()
	close(s.die)
}

// pushBytes a slice into buffer
func (s *Stream) pushBytes(p []byte) {
	s.bufferLock.Lock()
	s.buffer.Write(p)
	s.bufferLock.Unlock()
}

// recycleTokens transform remaining bytes to tokens(will truncate buffer)
func (s *Stream) recycleTokens() (n int) {
	s.bufferLock.Lock()
	n = s.buffer.Len()
	s.buffer.Reset()
	s.bufferLock.Unlock()
	return
}

// split large byte buffer into smaller frames, reference only
func (s *Stream) split(bts []byte, cmd byte, sid uint32) []Frame {
	var frames []Frame
	for len(bts) > int(s.frameSize) {
{'fraction_non_alphanumeric': 0.13668114026807626, 'fraction_numerical': 0.006796299792335284, 'mean_word_length': 3.4186822351959965, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 7, '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': '1481585', 'n_tokens_mistral': 2094, 'n_tokens_neox': 2006, 'n_words': 538}