<commit_before>package agent

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"net/url"
	"os"
	"strconv"
	"sync"

	"github.com/chrislusf/glow/driver/cmd"
	"github.com/chrislusf/glow/io/store"
	"github.com/chrislusf/glow/resource"
	"github.com/chrislusf/glow/resource/service_discovery/client"
	"github.com/chrislusf/glow/util"
	"github.com/golang/protobuf/proto"
)

type LiveDataStore struct {
	store           store.DataStore
	killHeartBeater chan bool
}

func NewLiveDataStore(s store.DataStore) *LiveDataStore {
	return &LiveDataStore{
		store:           s,
		killHeartBeater: make(chan bool, 1),
	}
}

func (ds *LiveDataStore) Destroy() {
	ds.killHeartBeater <- true
	ds.store.Destroy()
}

type AgentServerOption struct {
	Leader      *string
	Port        *int
	Dir         *string
	DataCenter  *string
	Rack        *string
	MaxExecutor *int
	MemoryMB    *int64
	CPULevel    *int
}

type AgentServer struct {
	Option                *AgentServerOption
	leader                string
	Port                  int
	name2Store            map[string]*LiveDataStore
	dir                   string
	name2StoreLock        sync.Mutex
	wg                    sync.WaitGroup
	l                     net.Listener
	computeResource       *resource.ComputeResource
	allocatedResource     *resource.ComputeResource
	allocatedResourceLock sync.Mutex
}

func NewAgentServer(option *AgentServerOption) *AgentServer {
	as := &AgentServer{
		Option:     option,
		leader:     *option.Leader,
		Port:       *option.Port,
		dir:        *option.Dir,
		name2Store: make(map[string]*LiveDataStore),
		computeResource: &resource.ComputeResource{
			CPUCount: *option.MaxExecutor,
			CPULevel: *option.CPULevel,
			MemoryMB: *option.MemoryMB,
		},
		allocatedResource: &resource.ComputeResource{},
	}

	err := as.Init()
	if err != nil {
		panic(err)
	}

	return as
}

// Start starts to listen on a port, returning the listening port
// r.Port can be pre-set or leave it as zero
// The actual port set to r.Port
func (r *AgentServer) Init() (err error) {
	r.l, err = net.Listen("tcp", ":"+strconv.Itoa(r.Port))
	if err != nil {
		log.Fatal(err)
	}

	r.Port = r.l.Addr().(*net.TCPAddr).Port
	fmt.Println("AgentServer starts on:", r.Port)
	return<commit_msg>add option to delete old .dat file on start<commit_after>package agent

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/chrislusf/glow/driver/cmd"
	"github.com/chrislusf/glow/io/store"
	"github.com/chrislusf/glow/resource"
	"github.com/chrislusf/glow/resource/service_discovery/client"
	"github.com/chrislusf/glow/util"
	"github.com/golang/protobuf/proto"
)

type LiveDataStore struct {
	store           store.DataStore
	killHeartBeater chan bool
}

func NewLiveDataStore(s store.DataStore) *LiveDataStore {
	return &LiveDataStore{
		store:           s,
		killHeartBeater: make(chan bool, 1),
	}
}

func (ds *LiveDataStore) Destroy() {
	ds.killHeartBeater <- true
	ds.store.Destroy()
}

type AgentServerOption struct {
	Leader       *string
	Port         *int
	Dir          *string
	DataCenter   *string
	Rack         *string
	MaxExecutor  *int
	MemoryMB     *int64
	CPULevel     *int
	CleanRestart *bool
}

type AgentServer struct {
	Option                *AgentServerOption
	leader                string
	Port                  int
	name2Store            map[string]*LiveDataStore
	dir                   string
	name2StoreLock        sync.Mutex
	wg                    sync.WaitGroup
	l                     net.Listener
	computeResource       *resource.ComputeResource
	allocatedResource     *resource.ComputeResource
	allocatedResourceLock sync.Mutex
}

func NewAgentServer(option *AgentServerOption) *AgentServer {
	as := &AgentServer{
		Option:     option,
		leader:     *option.Leader,
		Port:       *option.Port,
		dir:        *option.Dir,
		name2Store: make(map[string]*LiveDataStore),
		computeResource: &resource.ComputeResource{
			CPUCount: *option.MaxExecutor,
			CPULevel: *option.CPULevel,
			MemoryMB: *option.MemoryMB,
		},
		allocatedResource: &resource.ComputeResource{},
	}

	err := as.Init()
	if err != nil {
		panic(err)
	}

	return as
}

// Start starts to listen on a port, returning the listening port
// r.Port can be pre-set or leave it as zero
// The actual port set to r.Port
func (r *AgentServer) Init() (err error) {
	r.l, err = net.Listen("tcp", ":"+strconv.Itoa(r.Port))
	if err != nil {
		log.Fatal(err)
	}

	r.Port = r.l.Addr().(*net.TCPAddr).Port
	fmt.Println("AgentServer starts on:", r.Port)

	if *r.Option.CleanRestart {
		if fileInfos, err := ioutil.ReadDir(r.dir); err == nil {
			for _, fi := range fileInfos {
				name := fi.Name()
				if !fi.IsDir() && strings.HasSuffix(name, ".dat") {
					println("removing old dat file:", name)
					os.Remove(name)
				}
			}
		}
	}

	return
{'fraction_non_alphanumeric': 0.11532056619483763, 'fraction_numerical': 0.002497918401332223, 'mean_word_length': 2.816521048451152, 'pattern_counts': {'":': 2, '<': 5, '<?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': '7351732', 'n_tokens_mistral': 1774, 'n_tokens_neox': 1670, 'n_words': 373}