misc changes

projectdiscovery · Jul 1, 2024 · dba5a6d · dba5a6d
1 parent 3e08704
commit dba5a6d
Showing 1 changed file with 17 additions and 9 deletions.
diff --git a/batcher/batcher.go b/batcher/batcher.go
@@ -83,22 +83,28 @@ func New[T any](opts ...BatcherOption[T]) *Batcher[T] {
 // Append appends data to the batcher
 func (b *Batcher[T]) Append(d ...T) {
 	for _, item := range d {
-		sizeofItem := reflectutil.SizeOf(item)
-		currentSize := b.currentSize.Load()
+		var sizeofItem int
 
-		if b.maxSize > 0 && currentSize+int32(sizeofItem) > int32(b.maxSize) {
-			b.full <- true
-			b.incomingData <- item
-			b.currentSize.Add(int32(sizeofItem))
-			continue
+		if b.maxSize > 0 {
+			sizeofItem = reflectutil.SizeOf(item)
+			currentSize := b.currentSize.Load()
+
+			if currentSize+int32(sizeofItem) > int32(b.maxSize) {
+				b.full <- true
+				b.incomingData <- item
+				b.currentSize.Add(int32(sizeofItem))
+				continue
+			}
 		}
 
 		if !b.put(item) {
 			// will wait until space available
 			b.full <- true
 			b.incomingData <- item
 		}
-		b.currentSize.Add(int32(sizeofItem))
+		if b.maxSize > 0 {
+			b.currentSize.Add(int32(sizeofItem))
+		}
 	}
 }
 
@@ -176,7 +182,9 @@ func (b *Batcher[T]) doCallback() {
 	for item := range b.incomingData {
 		items[k] = item
 		k++
-		b.currentSize.Add(-int32(reflectutil.SizeOf(item)))
+		if b.maxSize > 0 {
+			b.currentSize.Add(-int32(reflectutil.SizeOf(item)))
+		}
 		if k >= n {
 			break
 		}