<commit_before>package termite

import (
	"net"
	"os"
	"log"
	"rpc"
	"sync"
)

// State associated with one master.
type Mirror struct {
	daemon         *WorkerDaemon
	fileServer     *rpc.Client
	fileServerConn net.Conn
	rpcConn        net.Conn
	rpcFs          *RpcFs
	writableRoot   string

	// key in WorkerDaemon's map.
	key string

	Waiting              int
	maxJobCount          int
	fuseFileSystemsMutex sync.Mutex
	// TODO - rename to unused FS.
	fuseFileSystems    []*WorkerFuseFs
	workingFileSystems map[*WorkerFuseFs]string
	shuttingDown       bool
	cond               sync.Cond
}

func NewMirror(daemon *WorkerDaemon, rpcConn, revConn net.Conn) *Mirror {
	log.Println("Mirror for", rpcConn, revConn)

	mirror := &Mirror{
		fileServerConn:     revConn,
		rpcConn:            rpcConn,
		fileServer:         rpc.NewClient(revConn),
		daemon:             daemon,
		workingFileSystems: make(map[*WorkerFuseFs]string),
	}
	mirror.rpcFs = NewRpcFs(mirror.fileServer, daemon.contentCache)
	mirror.rpcFs.localRoots = []string{"/lib", "/usr"}
	mirror.cond.L = &mirror.fuseFileSystemsMutex

	go mirror.serveRpc()
	return mirror
}

func (me *Mirror) DiscardFuse(wfs *WorkerFuseFs) {
	wfs.Stop()

	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()
	me.workingFileSystems[wfs] = "", false
	me.cond.Signal()
}

func (me *Mirror) serveRpc() {
	server := rpc.NewServer()
	server.Register(me)
	server.ServeConn(me.rpcConn)
	me.Shutdown()
}

func (me *Mirror) Shutdown() {
	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()
	if me.shuttingDown {
		return
	}
	me.shuttingDown = true
	me.fileServerConn.Close()
	me.fuseFileSystems = []*WorkerFuseFs{}

	for len(me.workingFileSystems) > 0 {
		// TODO - access is racy.
		me.maxJobCount = len(me.workingFileSystems)
		me.cond.Wait()
	}
	me.rpcConn.Close()

	go me.daemon.DropMirror(me)
}

func (me *Mirror) getWorkerFuseFs(name string) (f *WorkerFuseFs, err os.Error) {
	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()

	me.Waiting++
	for len(me.workingFileSystems) >= me.maxJobCount {
		me.cond.Wait()
	}
	me.Waiting--
	if me.shuttingDown {
		return nil, os.NewError("shutting down")
	}

	l := len(me.fuseFileSystems)
	if l > 0 {
		f = me.fuseFileSystems[l-1]
		me.fuseFileSystems = me.fuseFileSystems[:l-1]
	}
	if f == nil {
		f, err = me.newWorkerFuseFs()
	}
	me.workingFileSystems[f] = name
	return f, err
}

func (me *Mirror) Update(req *UpdateRequest, rep *UpdateResponse) os.Error {
	me.updateFiles(req.Files)
	return nil
}

func (me *Mirror) updateFiles(attrs []FileAttr) {
	me.updateFileSystems(attrs)
	me.rpcFs.updateFiles(attrs)
}

func (me *Mirror) updateFileSystems(attrs []FileAttr) {
	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()

	for _, fs := range me.fuseFileSystems {
		fs.update(attrs)
	}
	for fs, _ := range me.workingFileSystems {
		fs.update(attrs)
	}
}

func (me *Mirror) fetchFiles(files []FileAttr) {
	for _, f := range files {
		if f.Hash != nil && f.Content == nil {
			me.rpcFs.FetchHash(f.FileInfo.Size, f.Hash)
		}
	}
}

func (me *Mirror) Run(req *WorkRequest, rep *WorkReply) os.Error {
	me.updateFiles(req.Prefetch)
	go me.fetchFiles(req.Prefetch)
	task, err := me.newWorkerTask(req, rep)
	if err != nil {
		return err
	}

	err = task.Run()
	if err != nil {
		log.Println("Error", err)
		return err
	}

	updateReq := UpdateRequest{
		Files: rep.Files,
	}
	updateRep := UpdateResponse{}
	err = me.Update(&updateReq, &updateRep)
	if err != nil {
		return err
	}

	summary := rep
	// Trim output.

	summary.Stdout = trim(summary.Stdout)
	summary.Stderr = trim(summary.Stderr)

	log.Println("sending back", summary)
	return nil
}

const _DELETIONS = "DELETIONS"

func (me *Mirror) newWorkerFuseFs() (*WorkerFuseFs, os.Error) {
	return newWorkerFuseFs(me.daemon.tmpDir, me.rpcFs, me.writableRoot)
}

func (me *Mirror) newWorkerTask(req *WorkRequest, rep *WorkReply) (*WorkerTask, os.Error) {
	fuseFs, err := me.getWorkerFuseFs(req.Summary())
	if err != nil {
		return nil, err
	}
	var stdin net.Conn
	if req.StdinId != "" {
		stdin = me.daemon.pending.WaitConnection(req.StdinId)
	}<commit_msg>Polish.<commit_after>package termite

import (
	"net"
	"os"
	"log"
	"rpc"
	"sync"
)

// State associated with one master.
type Mirror struct {
	daemon         *WorkerDaemon
	fileServer     *rpc.Client
	fileServerConn net.Conn
	rpcConn        net.Conn
	rpcFs          *RpcFs
	writableRoot   string

	// key in WorkerDaemon's map.
	key string

	Waiting              int
	maxJobCount          int
	fuseFileSystemsMutex sync.Mutex
	// TODO - rename to unused FS.
	fuseFileSystems    []*WorkerFuseFs
	workingFileSystems map[*WorkerFuseFs]string
	shuttingDown       bool
	cond               sync.Cond
}

func NewMirror(daemon *WorkerDaemon, rpcConn, revConn net.Conn) *Mirror {
	log.Println("Mirror for", rpcConn, revConn)

	mirror := &Mirror{
		fileServerConn:     revConn,
		rpcConn:            rpcConn,
		fileServer:         rpc.NewClient(revConn),
		daemon:             daemon,
		workingFileSystems: make(map[*WorkerFuseFs]string),
	}
	mirror.rpcFs = NewRpcFs(mirror.fileServer, daemon.contentCache)
	mirror.rpcFs.localRoots = []string{"/lib", "/usr"}
	mirror.cond.L = &mirror.fuseFileSystemsMutex

	go mirror.serveRpc()
	return mirror
}

func (me *Mirror) DiscardFuse(wfs *WorkerFuseFs) {
	wfs.Stop()

	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()
	me.workingFileSystems[wfs] = "", false
	me.cond.Signal()
}

func (me *Mirror) serveRpc() {
	server := rpc.NewServer()
	server.Register(me)
	server.ServeConn(me.rpcConn)
	me.Shutdown()
}

func (me *Mirror) Shutdown() {
	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()
	if me.shuttingDown {
		return
	}
	me.shuttingDown = true
	me.fileServerConn.Close()
	me.fuseFileSystems = []*WorkerFuseFs{}

	for len(me.workingFileSystems) > 0 {
		// TODO - access is racy.
		me.maxJobCount = len(me.workingFileSystems)
		me.cond.Wait()
	}
	me.rpcConn.Close()

	go me.daemon.DropMirror(me)
}

func (me *Mirror) getWorkerFuseFs(name string) (f *WorkerFuseFs, err os.Error) {
	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()

	me.Waiting++
	for len(me.workingFileSystems) >= me.maxJobCount {
		me.cond.Wait()
	}
	me.Waiting--
	if me.shuttingDown {
		return nil, os.NewError("shutting down")
	}

	l := len(me.fuseFileSystems)
	if l > 0 {
		f = me.fuseFileSystems[l-1]
		me.fuseFileSystems = me.fuseFileSystems[:l-1]
	}
	if f == nil {
		f, err = me.newWorkerFuseFs()
	}
	me.workingFileSystems[f] = name
	return f, err
}

func (me *Mirror) Update(req *UpdateRequest, rep *UpdateResponse) os.Error {
	me.updateFiles(req.Files)
	return nil
}

func (me *Mirror) updateFiles(attrs []FileAttr) {
	me.rpcFs.updateFiles(attrs)
	
	me.fuseFileSystemsMutex.Lock()
	defer me.fuseFileSystemsMutex.Unlock()

	for _, fs := range me.fuseFileSystems {
		fs.update(attrs)
	}
	for fs, _ := range me.workingFileSystems {
		fs.update(attrs)
	}
}

func (me *Mirror) fetchFiles(files []FileAttr) {
	for _, f := range files {
		if f.Hash != nil && f.Content == nil {
			me.rpcFs.FetchHash(f.FileInfo.Size, f.Hash)
		}
	}
}

func (me *Mirror) Run(req *WorkRequest, rep *WorkReply) os.Error {
	me.updateFiles(req.Prefetch)
	go me.fetchFiles(req.Prefetch)
	task, err := me.newWorkerTask(req, rep)
	if err != nil {
		return err
	}

	err = task.Run()
	if err != nil {
		log.Println("Error", err)
		return err
	}

	me.updateFiles(rep.Files)

	summary := rep
	summary.Stdout = trim(summary.Stdout)
	summary.Stderr = trim(summary.Stderr)

	log.Println("sending back", summary)
	return nil
}

const _DELETIONS = "DELETIONS"

func (me *Mirror) newWorkerFuseFs() (*WorkerFuseFs, os.Error) {
	return newWorkerFuseFs(me.daemon.tmpDir, me.rpcFs, me.writableRoot)
}

func (me *Mirror) newWorkerTask(req *WorkRequest, rep *WorkReply) (*WorkerTask, os.Error) {
	fuseFs, err := me.getWorkerFuseFs(req.Summary())
	if err != nil {
		return nil, err
	}
	var stdin net.Conn
	if req.StdinId != "" {
		stdin = me.daemon.pending.WaitConnection(req.StdinId)
	}
{'fraction_non_alphanumeric': 0.13296948647978168, 'fraction_numerical': 0.0009923096005953857, 'mean_word_length': 3.7373678025851937, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '610339', 'n_tokens_mistral': 3284, 'n_tokens_neox': 3114, 'n_words': 633}