Fix FlatMap & BufferWithTimeOrCount

observable_operator.go

@@ -531,12 +531,22 @@ func (o *ObservableImpl) BufferWithTimeOrCount(timespan Duration, count int, opt
 
 		checkBuffer := func() {
 			mutex.Lock()
-			if len(buffer) != 0 {
-				if !Of(buffer).SendContext(ctx, next) {
-					mutex.Unlock()
-					return
+			switch {
+			case len(buffer) == 0:
+			case len(buffer) > count:
+				for {
+					if !Of(buffer[:count]).SendContext(ctx, next) {
+						break
+					}
+					buffer = buffer[count:]
+					if len(buffer) < count {
+						break
+					}
+				}
+			default:
+				if Of(buffer).SendContext(ctx, next) {
+					buffer = make([]interface{}, 0)
 				}
-				buffer = make([]interface{}, 0)
 			}
 			mutex.Unlock()
 		}
@@ -1111,44 +1121,40 @@ func (op *firstOrDefaultOperator) gatherNext(_ context.Context, _ Item, _ chan<-
 
 // FlatMap transforms the items emitted by an Observable into Observables, then flatten the emissions from those into a single Observable.
 func (o *ObservableImpl) FlatMap(apply ItemToObservable, opts ...Option) Observable {
-	f := func(ctx context.Context, next chan Item, option Option, opts ...Option) {
-		defer close(next)
-		observe := o.Observe(opts...)
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case item, ok := <-observe:
-				if !ok {
-					return
-				}
-				observe2 := apply(item).Observe(opts...)
-			loop2:
-				for {
-					select {
-					case <-ctx.Done():
-						return
-					case item, ok := <-observe2:
-						if !ok {
-							break loop2
-						}
-						if item.Error() {
-							item.SendContext(ctx, next)
-							if option.getErrorStrategy() == StopOnError {
-								return
-							}
-						} else {
-							if !item.SendContext(ctx, next) {
-								return
-							}
-						}
-					}
-				}
-			}
+	return observable(o, func() operator {
+		return &flatMapOperator{apply: apply}
+	}, false, true, opts...)
+}
+
+type flatMapOperator struct {
+	apply ItemToObservable
+}
+
+func (op *flatMapOperator) next(ctx context.Context, item Item, dst chan<- Item, operatorOptions operatorOptions) {
+	observe2 := op.apply(item).Observe()
+	for item := range observe2 {
+		if item.E != nil {
+			Error(item.E).SendContext(ctx, dst)
+			operatorOptions.stop()
+			return
 		}
+		Of(item.V).SendContext(ctx, dst)
 	}
+}
 
-	return customObservableOperator(f, opts...)
+func (op *flatMapOperator) err(ctx context.Context, item Item, dst chan<- Item, operatorOptions operatorOptions) {
+	op.next(ctx, item, dst, operatorOptions)
+}
+
+func (op *flatMapOperator) end(_ context.Context, _ chan<- Item) {
+}
+
+func (op *flatMapOperator) gatherNext(ctx context.Context, item Item, dst chan<- Item, _ operatorOptions) {
+	switch item.V.(type) {
+	case *flatMapOperator:
+		return
+	}
+	item.SendContext(ctx, dst)
 }
 
 // ForEach subscribes to the Observable and receives notifications for each element.

observable_operator_test.go

@@ -308,19 +308,14 @@ func Test_Observable_BufferWithTimeOrCount(t *testing.T) {
 	defer cancel()
 	ch := make(chan Item, 1)
 	obs := FromChannel(ch)
-	obs = obs.BufferWithTimeOrCount(WithDuration(30*time.Millisecond), 100)
+	obs = obs.BufferWithTimeOrCount(WithDuration(time.Second), 2)
 	go func() {
-		for i := 0; i < 10; i++ {
+		for i := 0; i < 5; i++ {
 			ch <- Of(i)
 		}
 		close(ch)
 	}()
-	Assert(ctx, t, obs, CustomPredicate(func(items []interface{}) error {
-		if len(items) == 0 {
-			return errors.New("items should not be nil")
-		}
-		return nil
-	}))
+	Assert(ctx, t, obs, HasItems([]interface{}{0, 1}, []interface{}{2, 3}, []interface{}{4}))
 }
 
 func Test_Observable_Contain(t *testing.T) {