zstd: Add ResetWithOptions to encoder/decoder (#1122)

## New API Methods

```
  func (*Encoder) ResetWithOptions(w io.Writer, opts ...EOption) error
  func (*Decoder) ResetWithOptions(r io.Reader, opts ...DOption) error

```

New Options

  // Encoder - clears dictionary
  `WithEncoderDictDelete() EOption`

  // Decoder - removes dicts by ID; no args clears all
  `WithDecoderDictDelete(ids ...uint32) DOption`

Option Reset Compatibility

  Can be changed with ResetWithOptions:
  - Encoder: WithEncoderCRC, WithEncoderPadding, WithZeroFrames, WithAllLitEntropyCompression, WithNoEntropyCompression, WithSingleSegment, WithEncoderDict, WithEncoderDictRaw, WithEncoderDictDelete
  - Decoder: WithDecoderMaxMemory, WithDecoderMaxWindow, WithDecoderDicts, WithDecoderDictRaw, WithDecoderDictDelete, WithDecodeAllCapLimit, IgnoreChecksum

  Cannot be changed with ResetWithOptions:
  - Encoder: WithEncoderConcurrency, WithWindowSize, WithEncoderLevel, WithLowerEncoderMem
  - Decoder: WithDecoderLowmem, WithDecoderConcurrency, WithDecodeBuffersBelow

Author: klauspost

zstd/decoder.go

@@ -39,9 +39,6 @@ type Decoder struct {
 
 	frame *frameDec
 
-	// Custom dictionaries.
-	dicts map[uint32]*dict
-
 	// streamWg is the waitgroup for all streams
 	streamWg sync.WaitGroup
 }
@@ -101,12 +98,10 @@ func NewReader(r io.Reader, opts ...DOption) (*Decoder, error) {
 		d.current.err = ErrDecoderNilInput
 	}
 
-	// Transfer option dicts.
-	d.dicts = make(map[uint32]*dict, len(d.o.dicts))
-	for _, dc := range d.o.dicts {
-		d.dicts[dc.id] = dc
+	// Initialize dict map if needed.
+	if d.o.dicts == nil {
+		d.o.dicts = make(map[uint32]*dict)
 	}
-	d.o.dicts = nil
 
 	// Create decoders
 	d.decoders = make(chan *blockDec, d.o.concurrent)
@@ -238,6 +233,21 @@ func (d *Decoder) Reset(r io.Reader) error {
 	return nil
 }
 
+// ResetWithOptions will reset the decoder and apply the given options
+// for the next stream or DecodeAll operation.
+// Options are applied on top of the existing options.
+// Some options cannot be changed on reset and will return an error.
+func (d *Decoder) ResetWithOptions(r io.Reader, opts ...DOption) error {
+	d.o.resetOpt = true
+	defer func() { d.o.resetOpt = false }()
+	for _, o := range opts {
+		if err := o(&d.o); err != nil {
+			return err
+		}
+	}
+	return d.Reset(r)
+}
+
 // drainOutput will drain the output until errEndOfStream is sent.
 func (d *Decoder) drainOutput() {
 	if d.current.cancel != nil {
@@ -930,7 +940,7 @@ decodeStream:
 }
 
 func (d *Decoder) setDict(frame *frameDec) (err error) {
-	dict, ok := d.dicts[frame.DictionaryID]
+	dict, ok := d.o.dicts[frame.DictionaryID]
 	if ok {
 		if debugDecoder {
 			println("setting dict", frame.DictionaryID)

zstd/decoder_options.go

@@ -20,10 +20,11 @@ type decoderOptions struct {
 	concurrent      int
 	maxDecodedSize  uint64
 	maxWindowSize   uint64
-	dicts           []*dict
+	dicts           map[uint32]*dict
 	ignoreChecksum  bool
 	limitToCap      bool
 	decodeBufsBelow int
+	resetOpt        bool
 }
 
 func (o *decoderOptions) setDefault() {
@@ -42,8 +43,15 @@ func (o *decoderOptions) setDefault() {
 
 // WithDecoderLowmem will set whether to use a lower amount of memory,
 // but possibly have to allocate more while running.
+// Cannot be changed with ResetWithOptions.
 func WithDecoderLowmem(b bool) DOption {
-	return func(o *decoderOptions) error { o.lowMem = b; return nil }
+	return func(o *decoderOptions) error {
+		if o.resetOpt && b != o.lowMem {
+			return errors.New("WithDecoderLowmem cannot be changed on Reset")
+		}
+		o.lowMem = b
+		return nil
+	}
 }
 
 // WithDecoderConcurrency sets the number of created decoders.
@@ -56,16 +64,20 @@ func WithDecoderLowmem(b bool) DOption {
 // The value supplied must be at least 0.
 // When a value of 0 is provided GOMAXPROCS will be used.
 // By default this will be set to 4 or GOMAXPROCS, whatever is lower.
+// Cannot be changed with ResetWithOptions.
 func WithDecoderConcurrency(n int) DOption {
 	return func(o *decoderOptions) error {
 		if n < 0 {
 			return errors.New("concurrency must be at least 0")
 		}
+		newVal := n
 		if n == 0 {
-			o.concurrent = runtime.GOMAXPROCS(0)
-		} else {
-			o.concurrent = n
+			newVal = runtime.GOMAXPROCS(0)
 		}
+		if o.resetOpt && newVal != o.concurrent {
+			return errors.New("WithDecoderConcurrency cannot be changed on Reset")
+		}
+		o.concurrent = newVal
 		return nil
 	}
 }
@@ -74,6 +86,7 @@ func WithDecoderConcurrency(n int) DOption {
 // non-streaming operations or maximum window size for streaming operations.
 // This can be used to control memory usage of potentially hostile content.
 // Maximum is 1 << 63 bytes. Default is 64GiB.
+// Can be changed with ResetWithOptions.
 func WithDecoderMaxMemory(n uint64) DOption {
 	return func(o *decoderOptions) error {
 		if n == 0 {
@@ -93,29 +106,37 @@ func WithDecoderMaxMemory(n uint64) DOption {
 // "zstd --train" from the Zstandard reference implementation.
 //
 // If several dictionaries with the same ID are provided, the last one will be used.
+// Can be changed with ResetWithOptions.
 //
 // [dictionary format]: https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#dictionary-format
 func WithDecoderDicts(dicts ...[]byte) DOption {
 	return func(o *decoderOptions) error {
+		if o.dicts == nil {
+			o.dicts = make(map[uint32]*dict)
+		}
 		for _, b := range dicts {
 			d, err := loadDict(b)
 			if err != nil {
 				return err
 			}
-			o.dicts = append(o.dicts, d)
+			o.dicts[d.id] = d
 		}
 		return nil
 	}
 }
 
 // WithDecoderDictRaw registers a dictionary that may be used by the decoder.
 // The slice content can be arbitrary data.
+// Can be changed with ResetWithOptions.
 func WithDecoderDictRaw(id uint32, content []byte) DOption {
 	return func(o *decoderOptions) error {
 		if bits.UintSize > 32 && uint(len(content)) > dictMaxLength {
 			return fmt.Errorf("dictionary of size %d > 2GiB too large", len(content))
 		}
-		o.dicts = append(o.dicts, &dict{id: id, content: content, offsets: [3]int{1, 4, 8}})
+		if o.dicts == nil {
+			o.dicts = make(map[uint32]*dict)
+		}
+		o.dicts[id] = &dict{id: id, content: content, offsets: [3]int{1, 4, 8}}
 		return nil
 	}
 }
@@ -125,6 +146,7 @@ func WithDecoderDictRaw(id uint32, content []byte) DOption {
 // The Decoder will likely allocate more memory based on the WithDecoderLowmem setting.
 // If WithDecoderMaxMemory is set to a lower value, that will be used.
 // Default is 512MB, Maximum is ~3.75 TB as per zstandard spec.
+// Can be changed with ResetWithOptions.
 func WithDecoderMaxWindow(size uint64) DOption {
 	return func(o *decoderOptions) error {
 		if size < MinWindowSize {
@@ -142,6 +164,7 @@ func WithDecoderMaxWindow(size uint64) DOption {
 // or any size set in WithDecoderMaxMemory.
 // This can be used to limit decoding to a specific maximum output size.
 // Disabled by default.
+// Can be changed with ResetWithOptions.
 func WithDecodeAllCapLimit(b bool) DOption {
 	return func(o *decoderOptions) error {
 		o.limitToCap = b
@@ -154,17 +177,37 @@ func WithDecodeAllCapLimit(b bool) DOption {
 // This typically uses less allocations but will have the full decompressed object in memory.
 // Note that DecodeAllCapLimit will disable this, as well as giving a size of 0 or less.
 // Default is 128KiB.
+// Cannot be changed with ResetWithOptions.
 func WithDecodeBuffersBelow(size int) DOption {
 	return func(o *decoderOptions) error {
+		if o.resetOpt && size != o.decodeBufsBelow {
+			return errors.New("WithDecodeBuffersBelow cannot be changed on Reset")
+		}
 		o.decodeBufsBelow = size
 		return nil
 	}
 }
 
 // IgnoreChecksum allows to forcibly ignore checksum checking.
+// Can be changed with ResetWithOptions.
 func IgnoreChecksum(b bool) DOption {
 	return func(o *decoderOptions) error {
 		o.ignoreChecksum = b
 		return nil
 	}
 }
+
+// WithDecoderDictDelete removes dictionaries by ID.
+// If no ids are passed, all dictionaries are deleted.
+// Should be used with ResetWithOptions.
+func WithDecoderDictDelete(ids ...uint32) DOption {
+	return func(o *decoderOptions) error {
+		if len(ids) == 0 {
+			clear(o.dicts)
+		}
+		for _, id := range ids {
+			delete(o.dicts, id)
+		}
+		return nil
+	}
+}

zstd/decoder_test.go

@@ -2213,3 +2213,129 @@ func TestWithDecodeAllCapLimit(t *testing.T) {
 		})
 	}
 }
+
+func TestDecoderResetWithOptions(t *testing.T) {
+	dec, err := NewReader(nil, WithDecoderConcurrency(1), WithDecoderLowmem(true))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer dec.Close()
+
+	// Test changing safe options
+	t.Run("change-maxmemory", func(t *testing.T) {
+		err := dec.ResetWithOptions(nil, WithDecoderMaxMemory(1<<20))
+		if err != nil {
+			t.Errorf("ResetWithOptions should allow changing maxMemory: %v", err)
+		}
+	})
+
+	t.Run("change-maxwindow", func(t *testing.T) {
+		err := dec.ResetWithOptions(nil, WithDecoderMaxWindow(1<<20))
+		if err != nil {
+			t.Errorf("ResetWithOptions should allow changing maxWindow: %v", err)
+		}
+	})
+
+	t.Run("change-ignorechecksum", func(t *testing.T) {
+		err := dec.ResetWithOptions(nil, IgnoreChecksum(true))
+		if err != nil {
+			t.Errorf("ResetWithOptions should allow changing ignoreChecksum: %v", err)
+		}
+	})
+
+	// Test error when changing unsafe options
+	t.Run("change-concurrency-error", func(t *testing.T) {
+		err := dec.ResetWithOptions(nil, WithDecoderConcurrency(4))
+		if err == nil {
+			t.Error("ResetWithOptions should error when changing concurrency")
+		}
+	})
+
+	t.Run("change-lowmem-error", func(t *testing.T) {
+		err := dec.ResetWithOptions(nil, WithDecoderLowmem(false))
+		if err == nil {
+			t.Error("ResetWithOptions should error when changing lowmem")
+		}
+	})
+
+	// Test same value is allowed
+	t.Run("same-concurrency-ok", func(t *testing.T) {
+		err := dec.ResetWithOptions(nil, WithDecoderConcurrency(1))
+		if err != nil {
+			t.Errorf("ResetWithOptions should allow same concurrency: %v", err)
+		}
+	})
+}
+
+func TestDecoderDictDelete(t *testing.T) {
+	dictContent := []byte("test dictionary content for decompression testing purposes")
+
+	dec, err := NewReader(nil, WithDecoderDictRaw(100, dictContent), WithDecoderDictRaw(200, dictContent))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer dec.Close()
+
+	if len(dec.o.dicts) != 2 {
+		t.Fatalf("expected 2 dicts, got %d", len(dec.o.dicts))
+	}
+
+	// Delete specific dict
+	err = dec.ResetWithOptions(nil, WithDecoderDictDelete(100))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(dec.o.dicts) != 1 {
+		t.Errorf("expected 1 dict after delete, got %d", len(dec.o.dicts))
+	}
+	if _, ok := dec.o.dicts[200]; !ok {
+		t.Error("dict 200 should still exist")
+	}
+
+	// Add another dict
+	err = dec.ResetWithOptions(nil, WithDecoderDictRaw(300, dictContent))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(dec.o.dicts) != 2 {
+		t.Errorf("expected 2 dicts after add, got %d", len(dec.o.dicts))
+	}
+
+	// Delete all dicts with no arguments
+	err = dec.ResetWithOptions(nil, WithDecoderDictDelete())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(dec.o.dicts) != 0 {
+		t.Errorf("expected 0 dicts after delete all, got %d", len(dec.o.dicts))
+	}
+}
+
+func TestDecoderDictDeleteMultiple(t *testing.T) {
+	dictContent := []byte("test dictionary content")
+
+	dec, err := NewReader(nil,
+		WithDecoderDictRaw(100, dictContent),
+		WithDecoderDictRaw(200, dictContent),
+		WithDecoderDictRaw(300, dictContent))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer dec.Close()
+
+	if len(dec.o.dicts) != 3 {
+		t.Fatalf("expected 3 dicts, got %d", len(dec.o.dicts))
+	}
+
+	// Delete multiple dicts in one call
+	err = dec.ResetWithOptions(nil, WithDecoderDictDelete(100, 300))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(dec.o.dicts) != 1 {
+		t.Errorf("expected 1 dict after delete, got %d", len(dec.o.dicts))
+	}
+	if _, ok := dec.o.dicts[200]; !ok {
+		t.Error("dict 200 should still exist")
+	}
+}

zstd/encoder.go

@@ -131,6 +131,22 @@ func (e *Encoder) Reset(w io.Writer) {
 	s.frameContentSize = 0
 }
 
+// ResetWithOptions will re-initialize the writer and apply the given options
+// as a new, independent stream.
+// Options are applied on top of the existing options.
+// Some options cannot be changed on reset and will return an error.
+func (e *Encoder) ResetWithOptions(w io.Writer, opts ...EOption) error {
+	e.o.resetOpt = true
+	defer func() { e.o.resetOpt = false }()
+	for _, o := range opts {
+		if err := o(&e.o); err != nil {
+			return err
+		}
+	}
+	e.Reset(w)
+	return nil
+}
+
 // ResetContentSize will reset and set a content size for the next stream.
 // If the bytes written does not match the size given an error will be returned
 // when calling Close().