X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/b6a55c2430033ee1a17b2266b529b425e641d19f..4d7567b7b3577b561dd064d397c10a9331c3ee16:/services/keep-balance/balance.go

diff --git a/services/keep-balance/balance.go b/services/keep-balance/balance.go
index bd3152fd88..9f814a20d3 100644
--- a/services/keep-balance/balance.go
+++ b/services/keep-balance/balance.go
@@ -78,11 +78,18 @@ func (bal *Balancer) Run(config Config, runOptions RunOptions) (nextRunOptions R
 		if err != nil {
 			return
 		}
+		defer lbFile.Close()
 		err = syscall.Flock(int(lbFile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB)
 		if err != nil {
 			return
 		}
-		defer os.Remove(tmpfn)
+		defer func() {
+			// Remove the tempfile only if we didn't get
+			// as far as successfully renaming it.
+			if lbFile != nil {
+				os.Remove(tmpfn)
+			}
+		}()
 		bal.lostBlocks = lbFile
 	} else {
 		bal.lostBlocks = ioutil.Discard
@@ -140,10 +147,7 @@ func (bal *Balancer) Run(config Config, runOptions RunOptions) (nextRunOptions R
 		if err != nil {
 			return
 		}
-		err = lbFile.Close()
-		if err != nil {
-			return
-		}
+		lbFile = nil
 	}
 	if runOptions.CommitPulls {
 		err = bal.CommitPulls(&config.Client)
@@ -446,7 +450,13 @@ func (bal *Balancer) addCollection(coll arvados.Collection) error {
 		repl = *coll.ReplicationDesired
 	}
 	debugf("%v: %d block x%d", coll.UUID, len(blkids), repl)
-	bal.BlockStateMap.IncreaseDesired(coll.StorageClassesDesired, repl, blkids)
+	// Pass pdh to IncreaseDesired only if LostBlocksFile is being
+	// written -- otherwise it's just a waste of memory.
+	pdh := ""
+	if bal.LostBlocksFile != "" {
+		pdh = coll.PortableDataHash
+	}
+	bal.BlockStateMap.IncreaseDesired(pdh, coll.StorageClassesDesired, repl, blkids)
 	return nil
 }
 
@@ -513,7 +523,7 @@ func (bal *Balancer) setupLookupTables() {
 				bal.mountsByClass["default"][mnt] = true
 				continue
 			}
-			for _, class := range mnt.StorageClasses {
+			for class := range mnt.StorageClasses {
 				if mbc := bal.mountsByClass[class]; mbc == nil {
 					bal.classes = append(bal.classes, class)
 					bal.mountsByClass[class] = map[*KeepMount]bool{mnt: true}
@@ -920,7 +930,11 @@ func (bal *Balancer) collectStatistics(results <-chan balanceResult) {
 			s.lost.replicas -= surplus
 			s.lost.blocks++
 			s.lost.bytes += bytes * int64(-surplus)
-			fmt.Fprintf(bal.lostBlocks, "%s\n", strings.SplitN(string(result.blkid), "+", 2)[0])
+			fmt.Fprintf(bal.lostBlocks, "%s", strings.SplitN(string(result.blkid), "+", 2)[0])
+			for pdh := range result.blk.Refs {
+				fmt.Fprintf(bal.lostBlocks, " %s", pdh)
+			}
+			fmt.Fprint(bal.lostBlocks, "\n")
 		case surplus < 0:
 			s.underrep.replicas -= surplus
 			s.underrep.blocks++