<commit_before>package ami

import (
	"bufio"
	"bytes"
	"io"
	"log"
	"net"
	"strings"
	"sync"
)

// Socket holds the socket client connection data.
type Socket struct {
	conn     net.Conn
	incoming chan string
	shutdown chan struct{}
	wg       sync.WaitGroup
}

// NewSocket provides a new socket client, connecting to a tcp server.
func NewSocket(address string) (*Socket, error) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		return nil, err
	}
	s := &Socket{
		conn:     conn,
		incoming: make(chan string, 32),
		shutdown: make(chan struct{}),
	}
	s.wg.Add(1)
	go s.run(conn)
	return s, nil
}

// Connected returns the socket status, true for connected,
// false for disconnected.
func (s *Socket) Connected() bool {
	return s.conn != nil
}

// Close closes socket connection.
func (s *Socket) Close() error {
	close(s.shutdown)
	s.wg.Wait()

	if s.conn != nil {
		return s.conn.Close()
	}
	return nil
}

// Send sends data to socket using fprintf format.
func (s *Socket) Send(message string) error {
	_, err := s.conn.Write([]byte(message))
	return err
}

// Recv receives a string from socket server.
func (s *Socket) Recv() (string, error) {
	var buffer bytes.Buffer
	for {
		select {
		case msg, ok := <-s.incoming:
			if !ok {
				continue
			}
			buffer.WriteString(msg)
			if strings.HasSuffix(buffer.String(), "\r\n") {
				return buffer.String(), nil
			}
		case <-s.shutdown:
			return buffer.String(), io.EOF
		}
	}
}

func (s *Socket) run(conn net.Conn) {
	defer s.wg.Done()

	reader := bufio.NewReader(conn)
	for {
		select {
		case <-s.shutdown:
			return
		default:
			msg, err := reader.ReadString('\n')
			if err != nil {
				s.Close()
				log.Fatalf("incoming message error: %v\n", err)
			}
			s.incoming <- msg
		}
	}
}<commit_msg>Fix waitgroup stuck<commit_after>package ami

import (
	"bufio"
	"bytes"
	"io"
	"log"
	"net"
	"strings"
	"sync"
)

// Socket holds the socket client connection data.
type Socket struct {
	conn     net.Conn
	incoming chan string
	shutdown chan struct{}
	wg       sync.WaitGroup
}

// NewSocket provides a new socket client, connecting to a tcp server.
func NewSocket(address string) (*Socket, error) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		return nil, err
	}
	s := &Socket{
		conn:     conn,
		incoming: make(chan string, 32),
		shutdown: make(chan struct{}),
	}
	s.wg.Add(1)
	go s.run(conn)
	return s, nil
}

// Connected returns the socket status, true for connected,
// false for disconnected.
func (s *Socket) Connected() bool {
	return s.conn != nil
}

// Close closes socket connection.
func (s *Socket) Close() error {
	close(s.shutdown)
	s.wg.Wait()

	if s.conn != nil {
		return s.conn.Close()
	}
	return nil
}

// Send sends data to socket using fprintf format.
func (s *Socket) Send(message string) error {
	_, err := s.conn.Write([]byte(message))
	return err
}

// Recv receives a string from socket server.
func (s *Socket) Recv() (string, error) {
	var buffer bytes.Buffer
	for {
		select {
		case msg, ok := <-s.incoming:
			if !ok {
				continue
			}
			buffer.WriteString(msg)
			if strings.HasSuffix(buffer.String(), "\r\n") {
				return buffer.String(), nil
			}
		case <-s.shutdown:
			return buffer.String(), io.EOF
		}
	}
}

func (s *Socket) run(conn net.Conn) {
	defer s.done()

	reader := bufio.NewReader(conn)
	for {
		select {
		case <-s.shutdown:
			return
		default:
			msg, err := reader.ReadString('\n')
			if err != nil {
				log.Printf("incoming message error: %v\n", err)
				return
			}
			s.incoming <- msg
		}
	}
}

func (s *Socket) done() {
	s.wg.Done()
	s.Close()
}
{'fraction_non_alphanumeric': 0.14480646059593427, 'fraction_numerical': 0.001670843776106934, 'mean_word_length': 3.0044593088071347, 'pattern_counts': {'":': 0, '<': 11, '<?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': '51605', 'n_tokens_mistral': 1432, 'n_tokens_neox': 1308, 'n_words': 396}