<commit_before>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/constants"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/srpc"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"github.com/Symantec/Dominator/sub/client"
	"runtime"
	"strings"
	"time"
)

func (sub *Sub) tryMakeBusy() bool {
	sub.busyMutex.Lock()
	defer sub.busyMutex.Unlock()
	if sub.busy {
		return false
	}
	sub.busy = true
	return true
}

func (sub *Sub) makeUnbusy() {
	sub.busyMutex.Lock()
	defer sub.busyMutex.Unlock()
	sub.busy = false
}

func (sub *Sub) connectAndPoll() {
	sub.status = statusConnecting
	hostname := strings.SplitN(sub.hostname, "*", 2)[0]
	address := fmt.Sprintf("%s:%d", hostname, constants.SubPortNumber)
	sub.lastConnectionStartTime = time.Now()
	srpcClient, err := srpc.DialHTTP("tcp", address)
	if err != nil {
		sub.status = statusFailedToConnect
		return
	}
	defer srpcClient.Close()
	sub.status = statusWaitingToPoll
	sub.lastConnectionSucceededTime = time.Now()
	sub.lastConnectDuration =
		sub.lastConnectionSucceededTime.Sub(sub.lastConnectionStartTime)
	connectDistribution.Add(sub.lastConnectDuration)
	sub.herd.pollSemaphore <- true
	sub.status = statusPolling
	sub.poll(srpcClient)
	<-sub.herd.pollSemaphore
}

func (sub *Sub) poll(srpcClient *srpc.Client) {
	var request subproto.PollRequest
	request.HaveGeneration = sub.generationCount
	var reply subproto.PollResponse
	sub.lastPollStartTime = time.Now()
	logger := sub.herd.logger
	if err := client.CallPoll(srpcClient, request, &reply); err != nil {
		sub.status = statusFailedToPoll
		logger.Printf("Error calling %s.Poll()\t%s\n", sub.hostname, err)
		return
	}
	sub.lastPollSucceededTime = time.Now()
	if reply.GenerationCount == 0 {
		sub.fileSystem = nil
		sub.objectCache = nil
		sub.generationCount = 0
	}
	if fs := reply.FileSystem; fs == nil {
		sub.lastShortPollDuration =
			sub.lastPollSucceededTime.Sub(sub.lastPollStartTime)
		shortPollDistribution.Add(sub.lastShortPollDuration)
	} else {
		if err := fs.RebuildInodePointers(); err != nil {
			sub.status = statusFailedToPoll
			logger.Printf("Error building pointers for: %s %s\n",
				sub.hostname, err)
			return
		}
		fs.BuildInodeToFilenamesTable()
		fs.BuildEntryMap()
		sub.fileSystem = fs
		sub.objectCache = reply.ObjectCache
		sub.generationCount = reply.GenerationCount
		sub.lastFullPollDuration =
			sub.lastPollSucceededTime.Sub(sub.lastPollStartTime)
		fullPollDistribution.Add(sub.lastFullPollDuration)
		logger.Printf("Polled: %s, GenerationCount=%d\n",
			sub.hostname, reply.GenerationCount)
	}
	if reply.FetchInProgress {
		sub.status = statusFetching
		return
	}
	if reply.UpdateInProgress {
		sub.status = statusUpdating
		return
	}
	if sub.generationCountAtLastSync == sub.generationCount {
		sub.status = statusSynced
		return
	}
	if sub.fileSystem == nil {
		sub.status = statusSubNotReady
		return
	}
	if sub.generationCountAtChangeStart == sub.generationCount {
		sub.status = statusWaitingForNextPoll
		return
	}
	if idle, status := sub.fetchMissingObjects(srpcClient,
		sub.requiredImage); !idle {
		sub.status = status
		sub.fileSystem = nil  // Mark memory for reclaim.
		sub.objectCache = nil // Mark memory for reclaim.
		runtime.GC()          // Reclaim now.
		return
	}
	sub.status = statusComputingUpdate
	if idle, status := sub.sendUpdate(srpcClient); !idle {
		sub.status = status
		sub.fileSystem = nil  // Mark memory for reclaim.
		sub.objectCache = nil // Mark memory for reclaim.
		runtime.GC()          // Reclaim now.
		return
	}
	if idle, status := sub.fetchMissingObjects(srpcClient,
		sub.plannedImage); !idle {
		if status != statusImageNotReady {
			sub.status = status
			sub.fileSystem = nil  // Mark memory for reclaim.
			sub.objectCache = nil // Mark memory for reclaim.
			runtime.GC()          // Reclaim now.
			return
		}
	}
	sub.status = statusSynced
	sub.cleanup(srpcClient, sub.plannedImage)
	sub.generationCountAtLastSync = sub.generationCount
	sub.fileSystem = nil  // Mark memory for reclaim.
	sub.objectCache = nil // Mark memory for reclaim.
	runtime.GC()          // Reclaim now.
}

// Returns true if all required objects are available.
func (sub *Sub) fetchMissingObjects(srpcClient *srpc.Client, imageName string) (
	bool, uint) {
	image := sub.herd.getImage(imageName)
	if image == nil {
		return false, statusImageNotReady
	}
	missingObjects := make(map[hash.Hash]struct{})
	for _, inode := range image.FileSystem.InodeTable {
		if inode, ok := inode.(*filesystem.RegularInode); ok {
			if inode.Size > 0 {
				missingObjects[inode.Hash] = struct{}{}<commit_msg>Move reclaim code into Sub.reclaim() method.<commit_after>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/constants"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/srpc"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"github.com/Symantec/Dominator/sub/client"
	"runtime"
	"strings"
	"time"
)

func (sub *Sub) tryMakeBusy() bool {
	sub.busyMutex.Lock()
	defer sub.busyMutex.Unlock()
	if sub.busy {
		return false
	}
	sub.busy = true
	return true
}

func (sub *Sub) makeUnbusy() {
	sub.busyMutex.Lock()
	defer sub.busyMutex.Unlock()
	sub.busy = false
}

func (sub *Sub) connectAndPoll() {
	sub.status = statusConnecting
	hostname := strings.SplitN(sub.hostname, "*", 2)[0]
	address := fmt.Sprintf("%s:%d", hostname, constants.SubPortNumber)
	sub.lastConnectionStartTime = time.Now()
	srpcClient, err := srpc.DialHTTP("tcp", address)
	if err != nil {
		sub.status = statusFailedToConnect
		return
	}
	defer srpcClient.Close()
	sub.status = statusWaitingToPoll
	sub.lastConnectionSucceededTime = time.Now()
	sub.lastConnectDuration =
		sub.lastConnectionSucceededTime.Sub(sub.lastConnectionStartTime)
	connectDistribution.Add(sub.lastConnectDuration)
	sub.herd.pollSemaphore <- true
	sub.status = statusPolling
	sub.poll(srpcClient)
	<-sub.herd.pollSemaphore
}

func (sub *Sub) poll(srpcClient *srpc.Client) {
	var request subproto.PollRequest
	request.HaveGeneration = sub.generationCount
	var reply subproto.PollResponse
	sub.lastPollStartTime = time.Now()
	logger := sub.herd.logger
	if err := client.CallPoll(srpcClient, request, &reply); err != nil {
		sub.status = statusFailedToPoll
		logger.Printf("Error calling %s.Poll()\t%s\n", sub.hostname, err)
		return
	}
	sub.lastPollSucceededTime = time.Now()
	if reply.GenerationCount == 0 {
		sub.reclaim()
		sub.generationCount = 0
	}
	if fs := reply.FileSystem; fs == nil {
		sub.lastShortPollDuration =
			sub.lastPollSucceededTime.Sub(sub.lastPollStartTime)
		shortPollDistribution.Add(sub.lastShortPollDuration)
	} else {
		if err := fs.RebuildInodePointers(); err != nil {
			sub.status = statusFailedToPoll
			logger.Printf("Error building pointers for: %s %s\n",
				sub.hostname, err)
			return
		}
		fs.BuildInodeToFilenamesTable()
		fs.BuildEntryMap()
		sub.fileSystem = fs
		sub.objectCache = reply.ObjectCache
		sub.generationCount = reply.GenerationCount
		sub.lastFullPollDuration =
			sub.lastPollSucceededTime.Sub(sub.lastPollStartTime)
		fullPollDistribution.Add(sub.lastFullPollDuration)
		logger.Printf("Polled: %s, GenerationCount=%d\n",
			sub.hostname, reply.GenerationCount)
	}
	if reply.FetchInProgress {
		sub.status = statusFetching
		return
	}
	if reply.UpdateInProgress {
		sub.status = statusUpdating
		return
	}
	if sub.generationCountAtLastSync == sub.generationCount {
		sub.status = statusSynced
		return
	}
	if sub.fileSystem == nil {
		sub.status = statusSubNotReady
		return
	}
	if sub.generationCountAtChangeStart == sub.generationCount {
		sub.status = statusWaitingForNextPoll
		return
	}
	if idle, status := sub.fetchMissingObjects(srpcClient,
		sub.requiredImage); !idle {
		sub.status = status
		sub.reclaim()
		return
	}
	sub.status = statusComputingUpdate
	if idle, status := sub.sendUpdate(srpcClient); !idle {
		sub.status = status
		sub.reclaim()
		return
	}
	if idle, status := sub.fetchMissingObjects(srpcClient,
		sub.plannedImage); !idle {
		if status != statusImageNotReady {
			sub.status = status
			sub.reclaim()
			return
		}
	}
	sub.status = statusSynced
	sub.cleanup(srpcClient, sub.plannedImage)
	sub.generationCountAtLastSync = sub.generationCount
	sub.reclaim()
}

func (sub *Sub) reclaim() {
	sub.fileSystem = nil  // Mark memory for reclaim.
	sub.objectCache = nil // Mark memory for reclaim.
	runtime.GC()          // Reclaim now.
}

// Returns true if all required objects are available.
func (sub *Sub) fetchMissingObjects(srpcClient *srpc.Client, imageName string) (
	bool, uint) {
	image := sub.herd.getImage(imageName)
	if image == nil {
		return false, statusImageNotReady
	}
	missingObjects := make(map[hash.Hash]struct{})
	for _, inode := range image.FileSystem.InodeTable {
		if inode, ok := inode.(*filesystem.RegularInode); ok {
			if inode.Size > 0 {
				missingObjects[inode.Hash] = struct{}{}
{'fraction_non_alphanumeric': 0.11954585930543188, 'fraction_numerical': 0.0011130899376669634, 'mean_word_length': 4.181660899653979, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, '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': '3986518', 'n_tokens_mistral': 3316, 'n_tokens_neox': 3147, 'n_words': 622}