<commit_before>package melody

import (
	"errors"
	"github.com/gorilla/websocket"
	"net/http"
	"time"
)

// A melody session.
type Session struct {
	Request *http.Request
	conn    *websocket.Conn
	output  chan *envelope
	melody  *Melody
}

func (s *Session) writeMessage(message *envelope) {
	select {
	case s.output <- message:
	default:
		s.melody.errorHandler(s, errors.New("Message buffer full"))
	}
}

func (s *Session) writeRaw(message *envelope) error {
	s.conn.SetWriteDeadline(time.Now().Add(s.melody.Config.WriteWait))
	err := s.conn.WriteMessage(message.t, message.msg)

	if err != nil {
		return err
	}

	if message.t == websocket.CloseMessage {
		err := s.conn.Close()

		if err != nil {
			return err
		}
	}

	return nil
}

func (s *Session) close() {
	s.writeRaw(&envelope{t: websocket.CloseMessage, msg: []byte{}})
}

func (s *Session) ping() {
	s.writeRaw(&envelope{t: websocket.PingMessage, msg: []byte{}})
}

func (s *Session) writePump() {
	defer s.conn.Close()

	ticker := time.NewTicker(s.melody.Config.PingPeriod)
	defer ticker.Stop()

loop:
	for {
		select {
		case msg, ok := <-s.output:
			if !ok {
				s.close()
				break loop
			}
			if err := s.writeRaw(msg); err != nil {
				s.melody.errorHandler(s, err)
				break loop
			}
		case <-ticker.C:
			s.ping()
		}
	}
}

func (s *Session) SetPongHandler(f func() error) {
	s.conn.SetPongHandler(func(string) error {
		s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
		return f()
	})
}

func (s *Session) readPump() {
	defer s.conn.Close()

	s.conn.SetReadLimit(s.melody.Config.MaxMessageSize)
	s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))

	s.SetPongHandler(func() error {
		return nil
	})

	for {
		t, message, err := s.conn.ReadMessage()

		if err != nil {
			s.melody.errorHandler(s, err)
			break
		}

		if t == websocket.TextMessage {
			s.melody.messageHandler(s, message)
		}

		if t == websocket.BinaryMessage {
			s.melody.messageHandlerBinary(s, message)<commit_msg>Use a mutex to ensure SetPongHandler atomicity<commit_after>package melody

import (
	"errors"
	"github.com/gorilla/websocket"
	"net/http"
	"sync"
	"time"
)

// A melody session.
type Session struct {
	Request *http.Request
	conn    *websocket.Conn
	mutex   sync.Mutex
	output  chan *envelope
	melody  *Melody
}

func (s *Session) writeMessage(message *envelope) {
	select {
	case s.output <- message:
	default:
		s.melody.errorHandler(s, errors.New("Message buffer full"))
	}
}

func (s *Session) writeRaw(message *envelope) error {
	s.conn.SetWriteDeadline(time.Now().Add(s.melody.Config.WriteWait))
	err := s.conn.WriteMessage(message.t, message.msg)

	if err != nil {
		return err
	}

	if message.t == websocket.CloseMessage {
		err := s.conn.Close()

		if err != nil {
			return err
		}
	}

	return nil
}

func (s *Session) close() {
	s.writeRaw(&envelope{t: websocket.CloseMessage, msg: []byte{}})
}

func (s *Session) ping() {
	s.writeRaw(&envelope{t: websocket.PingMessage, msg: []byte{}})
}

func (s *Session) writePump() {
	defer s.conn.Close()

	ticker := time.NewTicker(s.melody.Config.PingPeriod)
	defer ticker.Stop()

loop:
	for {
		select {
		case msg, ok := <-s.output:
			if !ok {
				s.close()
				break loop
			}
			if err := s.writeRaw(msg); err != nil {
				s.melody.errorHandler(s, err)
				break loop
			}
		case <-ticker.C:
			s.ping()
		}
	}
}

func (s *Session) SetPongHandler(f func() error) {
	s.mutex.Lock()
	{
		s.conn.SetPongHandler(func(string) error {
			s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
			return f()
		})
	}
	s.mutex.Unlock()
}

func (s *Session) readPump() {
	defer s.conn.Close()

	s.conn.SetReadLimit(s.melody.Config.MaxMessageSize)
	s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))

	s.mutex.Lock()
	{
		s.conn.SetPongHandler(func(string) error {
			s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
			return nil
		})
	}
	s.mutex.Unlock()

	for {
		t, message, err := s.conn.ReadMessage()

		if err != nil {
			s.melody.errorHandler(s, err)
			break
		}

		if t == websocket.TextMessage {
			s.melody.messageHandler(s, message)
		}

		if t == websocket.BinaryMessage {
			s.melody.messageHandlerBinary(s, message)
{'fraction_non_alphanumeric': 0.15766944114149822, 'fraction_numerical': 0.0, 'mean_word_length': 3.4935897435897436, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 3, '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': '2295018', 'n_tokens_mistral': 1791, 'n_tokens_neox': 1639, 'n_words': 327}