Update totalBytes and writePos when writing number of messages, and resolve the writeMessages off by 1 error.

Author: kev1n80

diskqueue.go

@@ -435,12 +435,14 @@ func (d *diskQueue) writeOne(data []byte) error {
 	totalBytes := int64(4 + dataLen)
 
 	// check if we reached the file size limit with this message
-	if d.diskLimitFeatIsOn && d.writePos+totalBytes+8 >= d.maxBytesPerFile {
+	if d.diskLimitFeatIsOn && d.writePos+totalBytes >= d.maxBytesPerFile {
 		// write number of messages in binary to file
-		err = binary.Write(&d.writeBuf, binary.BigEndian, d.writeMessages)
+		err = binary.Write(&d.writeBuf, binary.BigEndian, d.writeMessages+1)
 		if err != nil {
 			return err
 		}
+
+		totalBytes += 8
 	}
 
 	// only write to the file once
@@ -454,14 +456,11 @@ func (d *diskQueue) writeOne(data []byte) error {
 	d.writePos += totalBytes
 	d.depth += 1
 
-	fileSize := d.writePos
-
 	if d.diskLimitFeatIsOn {
-		// save space for the number of messages in this file
-		fileSize += 8
+		d.writeMessages += 1
 	}
 
-	if fileSize >= d.maxBytesPerFile {
+	if d.writePos >= d.maxBytesPerFile {
 		if d.readFileNum == d.writeFileNum {
 			d.maxBytesPerFileRead = d.writePos
 		}
@@ -480,8 +479,6 @@ func (d *diskQueue) writeOne(data []byte) error {
 			d.writeFile.Close()
 			d.writeFile = nil
 		}
-	} else {
-		d.writeMessages += 1
 	}
 
 	return err

diskqueue_test.go

@@ -259,22 +259,22 @@ type md struct {
 	writePos      int64
 }
 
-func readMetaDataFile(fileName string, retried int, withDiskSpaceFeat bool) md {
+func readMetaDataFile(fileName string, retried int, diskLimitFeatIsOn bool) md {
 	f, err := os.OpenFile(fileName, os.O_RDONLY, 0600)
 	if err != nil {
 		// provide a simple retry that results in up to
 		// another 500ms for the file to be written.
 		if retried < 9 {
 			retried++
 			time.Sleep(50 * time.Millisecond)
-			return readMetaDataFile(fileName, retried, withDiskSpaceFeat)
+			return readMetaDataFile(fileName, retried, diskLimitFeatIsOn)
 		}
 		panic(err)
 	}
 	defer f.Close()
 
 	var ret md
-	if withDiskSpaceFeat {
+	if diskLimitFeatIsOn {
 		_, err = fmt.Fscanf(f, "%d\n%d,%d,%d\n%d,%d,%d\n",
 			&ret.depth,
 			&ret.readFileNum, &ret.readMessages, &ret.readPos,