few changes

Author: OneOfOne

cmap.go

@@ -1,5 +1,7 @@
 package cmap
 
+import "context"
+
 type (
 
 	// KT is the KeyType of the map, used for generating specialized versions.
@@ -129,7 +131,7 @@ func (cm *CMap) Keys() []KT {
 
 // ForEach loops over all the key/values in all the shards in order.
 // You can break early by returning an error.
-// It is safe to change the map during this call.
+// It **is** safe to modify the map while using this iterator, however it uses more memory and is slightly slower.
 func (cm *CMap) ForEach(fn func(key KT, val VT) error) error {
 	for i := range cm.shards {
 		if err := cm.shards[i].ForEach(fn); err != nil {
@@ -144,10 +146,10 @@ func (cm *CMap) ForEach(fn func(key KT, val VT) error) error {
 
 // ForEachLocked loops over all the key/values in the map.
 // You can break early by returning an error.
-// It is **NOT** safe to change the map during this call.
+// It is **NOT* safe to modify the map while using this iterator.
 func (cm *CMap) ForEachLocked(fn func(key KT, val VT) error) error {
 	for i := range cm.shards {
-		if err := cm.shards[i].ForEach(fn); err != nil {
+		if err := cm.shards[i].ForEachLocked(fn); err != nil {
 			if err == Break {
 				return nil
 			}
@@ -158,49 +160,45 @@ func (cm *CMap) ForEachLocked(fn func(key KT, val VT) error) error {
 }
 
 // Iter returns a channel to be used in for range.
-// **Warning** that breaking early will leak up to cm.NumShards() goroutines, use IterWithCancel if you intend to break early.
-// It is safe to modify the map while using the iterator.
-func (cm *CMap) Iter(buffer int) <-chan *KV {
-	ch, _ := cm.IterWithCancel(buffer)
+// Use `context.WithCancel` if you intend to break early or goroutines will leak.
+// It **is** safe to modify the map while using this iterator, however it uses more memory and is slightly slower.
+func (cm *CMap) Iter(ctx context.Context, buffer int) <-chan *KV {
+	ch := make(chan *KV, buffer)
+	go func() {
+		cm.iterContext(ctx, ch, false)
+		close(ch)
+	}()
 	return ch
 }
 
-// IterWithCancel returns a channel to be used in for range and
-// a cancelFn that can be called at any time to cleanly exit early.
-// Note that cancelFn will block until all the writers are notified.
-// It is safe to modify the map while using the iterator.
-func (cm *CMap) IterWithCancel(buffer int) (kvChan <-chan *KV, cancelFn func()) {
-	var (
-		ch       = make(chan *KV, buffer)
-		cancelCh = make(chan struct{})
-	)
+// IterLocked returns a channel to be used in for range.
+// Use `context.WithCancel` if you intend to break early or goroutines will leak and map access will deadlock.
+// It is **NOT* safe to modify the map while using this iterator.
+func (cm *CMap) IterLocked(ctx context.Context, buffer int) <-chan *KV {
+	ch := make(chan *KV, buffer)
+	go func() {
+		cm.iterContext(ctx, ch, false)
+		close(ch)
+	}()
+	return ch
+}
 
-	kvChan, cancelFn = ch, func() {
+func (cm *CMap) iterContext(ctx context.Context, ch chan<- *KV, locked bool) {
+	fn := func(k KT, v VT) error {
 		select {
-		case <-cancelCh:
-		default:
-			close(cancelCh)
-			for range ch {
-			}
+		case <-ctx.Done():
+			return Break
+		case ch <- &KV{k, v}:
+			return nil
 		}
 	}
-
-	go func() {
-		for i := range cm.shards {
-			cm.shards[i].ForEach(func(k KT, v VT) error {
-				select {
-				case <-cancelCh:
-					return Break
-				case ch <- &KV{k, v}:
-					return nil
-				}
-			})
+	for i := range cm.shards {
+		if locked {
+			cm.shards[i].ForEachLocked(fn)
+		} else {
+			cm.shards[i].ForEach(fn)
 		}
-		close(ch)
-		cancelFn()
-	}()
-
-	return
+	}
 }
 
 // Len returns the number of elements in the map.

cmap_go18_test.go

@@ -1,6 +1,7 @@
 package cmap_test
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"sync"
@@ -22,9 +23,10 @@ func TestIter(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		cm.Set(i, i)
 	}
-	ch, cancel := cm.IterWithCancel(0)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	i := 0
-	for kv := range ch {
+	for kv := range cm.IterLocked(ctx, 1) {
 		t.Logf("%d: %+v", i, kv)
 		if i++; i > 10 {
 			cancel()

stringcmap/cmap.go

@@ -1,6 +1,8 @@
 package stringcmap
 
 import (
+	"context"
+
 	"github.com/OneOfOne/cmap"
 )
 
@@ -145,49 +147,28 @@ type KV struct {
 }
 
 // Iter returns a channel to be used in for range.
-// **Warning** that breaking early will leak up to cm.NumShards() goroutines, use IterWithCancel if you intend to break early.
-// It is safe to modify the map while using the iterator.
-func (cm *CMap) Iter(buffer int) <-chan *KV {
-	ch, _ := cm.IterWithCancel(buffer)
+// **Warning** breaking early will leak up to cm.NumShards() goroutines, use IterWithCancel if you intend to break early.
+// It is safe to modify the map while using this iterator.
+func (cm *CMap) Iter(ctx context.Context, buffer int) <-chan *KV {
+	ch := make(chan *KV, buffer)
+	go func() {
+		cm.iterContext(ctx, ch)
+		close(ch)
+	}()
 	return ch
 }
 
-// IterWithCancel returns a channel to be used in for range and
-// a cancelFn that can be called at any time to cleanly exit early.
-// Note that cancelFn will block until all the writers are notified.
-// It is safe to modify the map while using the iterator.
-func (cm *CMap) IterWithCancel(buffer int) (kvChan <-chan *KV, cancelFn func()) {
-	var (
-		ch       = make(chan *KV, buffer)
-		cancelCh = make(chan struct{})
-	)
-
-	kvChan, cancelFn = ch, func() {
-		select {
-		case <-cancelCh:
-		default:
-			close(cancelCh)
-			for range ch {
+func (cm *CMap) iterContext(ctx context.Context, ch chan<- *KV) {
+	for i := range cm.shards {
+		cm.shards[i].ForEach(func(k string, v interface{}) error {
+			select {
+			case <-ctx.Done():
+				return cmap.Break
+			case ch <- &KV{k, v}:
+				return nil
 			}
-		}
+		})
 	}
-
-	go func() {
-		for i := range cm.shards {
-			cm.shards[i].ForEach(func(k string, v interface{}) error {
-				select {
-				case <-cancelCh:
-					return cmap.Break
-				case ch <- &KV{k, v}:
-					return nil
-				}
-			})
-		}
-		close(ch)
-		cancelFn()
-	}()
-
-	return
 }
 
 // Len returns the number of elements in the map.

stringcmap/json_support.go

@@ -1,8 +1,10 @@
 package stringcmap
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/json"
+	"io"
 )
 
 // WithJSON returns a MapWithJSON with the specific unmarshal func.
@@ -22,30 +24,82 @@ type MapWithJSON struct {
 	UnmarshalValueFn func(raw json.RawMessage) (interface{}, error)
 }
 
-// UnmarshalJSON implements json.Unmarshaler
-func (mwj *MapWithJSON) UnmarshalJSON(p []byte) error {
-	if mwj.CMap == nil {
-		mwj.CMap = New()
+// WriteTo implements io.WriterTo, outputs the map as a json object.
+func (mjw *MapWithJSON) WriteTo(w io.Writer) (n int64, err error) {
+	var buf writerWithBytes
+	switch w := w.(type) {
+	case writerWithBytes:
+		buf = w
+	default:
+		buf = bufio.NewWriter(w)
 	}
 
-	if mwj.UnmarshalValueFn == nil {
-		var in map[string]interface{}
-		if err := json.Unmarshal(p, &in); err != nil {
+	buf.WriteByte('{')
+
+	if err := mjw.ForEach(func(key string, val interface{}) error {
+		vj, err := json.Marshal(val)
+		if err != nil {
 			return err
 		}
-
-		for k, v := range in {
-			mwj.Set(k, v)
+		if n > 0 {
+			buf.WriteByte(',')
 		}
-
+		kj, _ := json.Marshal(key)
+		kn, _ := buf.Write(kj)
+		buf.WriteByte(':')
+		vn, _ := buf.Write(vj)
+		n += int64(kn + vn + 1)
 		return nil
+	}); err != nil {
+		return 0, err
+	}
+
+	buf.WriteByte('}')
+	n += 2 // {}
+
+	if buf, ok := buf.(flusher); ok {
+		err = buf.Flush()
+	}
+
+	return
+}
+
+// UnmarshalJSON implements json.Unmarshaler
+func (mwj *MapWithJSON) UnmarshalJSON(p []byte) error {
+	if mwj.UnmarshalValueFn != nil {
+		return mwj.unmarshalJSONTyped(p)
+	}
+
+	return mwj.unmarshalJSON(p)
+}
+
+func (mwj *MapWithJSON) unmarshalJSON(p []byte) error {
+	var in map[string]interface{}
+	if err := json.Unmarshal(p, &in); err != nil {
+		return err
 	}
 
+	if len(in) > 0 && mwj.CMap == nil {
+		mwj.CMap = New()
+	}
+
+	for k, v := range in {
+		mwj.Set(k, v)
+	}
+
+	return nil
+}
+
+func (mwj *MapWithJSON) unmarshalJSONTyped(p []byte) error {
 	var in map[string]json.RawMessage
 	if err := json.Unmarshal(p, &in); err != nil {
 		return err
 	}
 
+	if len(in) > 0 && mwj.CMap == nil {
+		mwj.CMap = New()
+	}
+
 	for k, rj := range in {
 		v, err := mwj.UnmarshalValueFn(rj)
 		if err != nil {
@@ -60,26 +114,18 @@ func (mwj *MapWithJSON) UnmarshalJSON(p []byte) error {
 
 // MarshalJSON implements json.Marshaler.
 func (cm *CMap) MarshalJSON() ([]byte, error) {
-	buf := bytes.NewBufferString("{")
-	if err := cm.ForEach(func(key string, val interface{}) error {
-		vj, err := json.Marshal(val)
-		if err != nil {
-			return err
-		}
-		kj, _ := json.Marshal(key)
-		buf.Write(kj)
-		buf.WriteByte(':')
-		buf.Write(vj)
-		buf.WriteByte(',')
-		return nil
-	}); err != nil {
+	var buf bytes.Buffer
+	if _, err := cm.WithJSON(nil).WriteTo(&buf); err != nil {
 		return nil, err
 	}
-	out := buf.Bytes()
-	if out[len(out)-1] == ',' {
-		out[len(out)-1] = '}'
-	} else {
-		out = append(out, '}')
-	}
-	return out, nil
+	return buf.Bytes(), nil
+}
+
+type writerWithBytes interface {
+	io.Writer
+	io.ByteWriter
+}
+
+type flusher interface {
+	Flush() error
 }

stringcmap/json_support_test.go

@@ -1,6 +1,7 @@
 package stringcmap
 
 import (
+	"context"
 	"encoding/json"
 	"reflect"
 	"sort"
@@ -56,7 +57,7 @@ func TestJSONType(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	for kv := range cm.Iter(0) {
+	for kv := range cm.Iter(context.Background(), 0) {
 		if v, ok := kv.Value.(uint64); !ok || kv.Key != "1" || v != 1 {
 			t.Fatalf("bad kv: %#+v", kv)
 		}