feat: add generic filter iterator and use for fitlering

Author: 2color

routing/http/server/filters.go

@@ -1,11 +1,13 @@
 package server
 
 import (
+	"errors"
 	"reflect"
 	"slices"
 	"strings"
 
 	"github.com/ipfs/boxo/routing/http/types"
+	"github.com/ipfs/boxo/routing/http/types/iter"
 	"github.com/multiformats/go-multiaddr"
 )
 
@@ -18,18 +20,68 @@ func parseFilter(param string) []string {
 	return strings.Split(strings.ToLower(param), ",")
 }
 
-func filterProviders(providers []types.Record, filterAddrs, filterProtocols []string) []types.Record {
+// applyFiltersToIter applies the filters to the given iterator and returns a new iterator.
+func applyFiltersToIter(recordsIter iter.ResultIter[types.Record], filterAddrs, filterProtocols []string) iter.ResultIter[types.Record] {
+	mappedIter := iter.Map(recordsIter, func(v iter.Result[types.Record]) iter.Result[types.Record] {
+		if v.Err != nil || v.Val == nil {
+			return v
+		}
+
+		switch v.Val.GetSchema() {
+		case types.SchemaPeer:
+			record, ok := v.Val.(*types.PeerRecord)
+			if !ok {
+				logger.Errorw("problem casting find providers record", "Schema", v.Val.GetSchema(), "Type", reflect.TypeOf(v).String())
+				// TODO: Do we want to let failed type assertions to pass through?
+				return v
+			}
+
+			record = applyFilters(record, filterAddrs, filterProtocols)
+			if record == nil {
+				return iter.Result[types.Record]{Err: errors.New("record is nil")}
+			}
+			v.Val = record
+
+		//lint:ignore SA1019 // ignore staticcheck
+		case types.SchemaBitswap:
+			//lint:ignore SA1019 // ignore staticcheck
+			record, ok := v.Val.(*types.BitswapRecord)
+			if !ok {
+				logger.Errorw("problem casting find providers record", "Schema", v.Val.GetSchema(), "Type", reflect.TypeOf(v).String())
+				// TODO: Do we want to let failed type assertions to pass through?
+				return v
+			}
+			peerRecord := types.FromBitswapRecord(record)
+			peerRecord = applyFilters(peerRecord, filterAddrs, filterProtocols)
+			if peerRecord == nil {
+				return iter.Result[types.Record]{Err: errors.New("record is nil")}
+			}
+			v.Val = peerRecord
+		}
+		return v
+	})
+
+	// filter out nil results and errors
+	filteredIter := iter.Filter(mappedIter, func(v iter.Result[types.Record]) bool {
+		return v.Err == nil && v.Val != nil
+	})
+
+	return filteredIter
+}
+
+func filterRecords(records []types.Record, filterAddrs, filterProtocols []string) []types.Record {
 	if len(filterAddrs) == 0 && len(filterProtocols) == 0 {
-		return providers
+		return records
 	}
 
-	filtered := make([]types.Record, 0, len(providers))
+	filtered := make([]types.Record, 0, len(records))
 
-	for _, provider := range providers {
-		if schema := provider.GetSchema(); schema == types.SchemaPeer {
-			peer, ok := provider.(*types.PeerRecord)
+	for _, record := range records {
+		// TODO: Handle SchemaBitswap
+		if schema := record.GetSchema(); schema == types.SchemaPeer {
+			peer, ok := record.(*types.PeerRecord)
 			if !ok {
-				logger.Errorw("problem casting find providers result", "Schema", provider.GetSchema(), "Type", reflect.TypeOf(provider).String())
+				logger.Errorw("problem casting find providers result", "Schema", record.GetSchema(), "Type", reflect.TypeOf(record).String())
 				// if the type assertion fails, we exlude record from results
 				continue
 			}
@@ -42,7 +94,7 @@ func filterProviders(providers []types.Record, filterAddrs, filterProtocols []st
 
 		} else {
 			// Will we ever encounter the SchemaBitswap type? Evidence seems to suggest that no longer
-			logger.Errorw("encountered unknown provider schema", "Schema", provider.GetSchema(), "Type", reflect.TypeOf(provider).String())
+			logger.Errorw("encountered unknown provider schema", "Schema", record.GetSchema(), "Type", reflect.TypeOf(record).String())
 		}
 	}
 	return filtered
@@ -51,6 +103,10 @@ func filterProviders(providers []types.Record, filterAddrs, filterProtocols []st
 // Applies the filters. Returns nil if the provider does not pass the protocols filter
 // The address filter is more complicated because it potentially modifies the Addrs slice.
 func applyFilters(provider *types.PeerRecord, filterAddrs, filterProtocols []string) *types.PeerRecord {
+	if len(filterAddrs) == 0 && len(filterProtocols) == 0 {
+		return provider
+	}
+
 	if !applyProtocolFilter(provider.Protocols, filterProtocols) {
 		// If the provider doesn't match any of the passed protocols, the provider is omitted from the response.
 		return nil

routing/http/server/server.go

@@ -235,84 +235,21 @@ func (s *server) findProviders(w http.ResponseWriter, httpReq *http.Request) {
 func (s *server) findProvidersJSON(w http.ResponseWriter, provIter iter.ResultIter[types.Record], filterAddrs, filterProtocols []string) {
 	defer provIter.Close()
 
-	providers, err := iter.ReadAllResults(provIter)
+	filteredIter := applyFiltersToIter(provIter, filterAddrs, filterProtocols)
+	providers, err := iter.ReadAllResults(filteredIter)
 	if err != nil {
 		writeErr(w, "FindProviders", http.StatusInternalServerError, fmt.Errorf("delegate error: %w", err))
 		return
 	}
 
-	filteredProviders := filterProviders(providers, filterAddrs, filterProtocols)
-
 	writeJSONResult(w, "FindProviders", jsontypes.ProvidersResponse{
-		Providers: filteredProviders,
+		Providers: providers,
 	})
 }
 func (s *server) findProvidersNDJSON(w http.ResponseWriter, provIter iter.ResultIter[types.Record], filterAddrs, filterProtocols []string) {
-	defer provIter.Close()
-
-	w.Header().Set("Content-Type", mediaTypeNDJSON)
-	w.Header().Add("Vary", "Accept")
-	w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
-
-	hasResults := false
-	for provIter.Next() {
-		res := provIter.Val()
-		if res.Err != nil {
-			logger.Errorw("ndjson iterator error", "Error", res.Err)
-			return
-		}
-
-		// handle filtering per record as we iterate
-		if len(filterAddrs) > 0 || len(filterProtocols) > 0 {
-			switch v := res.Val.(type) {
-			case *types.PeerRecord:
-				record := applyFilters(v, filterAddrs, filterProtocols)
-				if record == nil {
-					// if the record is nil, we skip it
-					continue
-				}
-				res.Val = record
-			default:
-				logger.Warn("unexpected type for res.Val, expected types.PeerRecord")
-				continue
-			}
-		}
-
-		// don't use an encoder because we can't easily differentiate writer errors from encoding errors
-		b, err := drjson.MarshalJSONBytes(res.Val)
-		if err != nil {
-			logger.Errorw("ndjson marshal error", "Error", err)
-			return
-		}
-
-		if !hasResults {
-			hasResults = true
-			// There's results, cache useful result for longer
-			setCacheControl(w, maxAgeWithResults, maxStale)
-		}
+	filteredIter := applyFiltersToIter(provIter, filterAddrs, filterProtocols)
 
-		_, err = w.Write(b)
-		if err != nil {
-			logger.Warn("ndjson write error", "Error", err)
-			return
-		}
-
-		_, err = w.Write([]byte{'\n'})
-		if err != nil {
-			logger.Warn("ndjson write error", "Error", err)
-			return
-		}
-
-		if f, ok := w.(http.Flusher); ok {
-			f.Flush()
-		}
-	}
-
-	if !hasResults {
-		// There weren't results, cache for shorter and send 404
-		setCacheControl(w, maxAgeWithoutResults, maxStale)
-		w.WriteHeader(http.StatusNotFound)
-	}
+	writeResultsIterNDJSON(w, filteredIter)
 }
 
 func (s *server) findPeers(w http.ResponseWriter, r *http.Request) {

routing/http/types/iter/filter.go

@@ -0,0 +1,43 @@
+package iter
+
+// Filter returns an iterator that filters out values that don't satisfy the predicate f.
+func Filter[T any](iter Iter[T], f func(t T) bool) *FilterIter[T] {
+	return &FilterIter[T]{iter: iter, f: f}
+}
+
+type FilterIter[T any] struct {
+	iter Iter[T]
+	f    func(T) bool
+
+	done bool
+	val  T
+}
+
+func (f *FilterIter[T]) Next() bool {
+	if f.done {
+		return false
+	}
+
+	ok := f.iter.Next()
+	f.done = !ok
+
+	if f.done {
+		return false
+	}
+
+	f.val = f.iter.Val()
+
+	if f.f(f.val) {
+		return true
+	}
+
+	return f.Next()
+}
+
+func (f *FilterIter[T]) Val() T {
+	return f.val
+}
+
+func (f *FilterIter[T]) Close() error {
+	return f.iter.Close()
+}

routing/http/types/iter/filter_test.go

@@ -0,0 +1,41 @@
+package iter
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFilter(t *testing.T) {
+	for _, c := range []struct {
+		input      Iter[int]
+		f          func(int) bool
+		expResults []int
+	}{
+		{
+			input:      FromSlice([]int{1, 2, 3, 4}),
+			f:          func(i int) bool { return i%2 == 0 },
+			expResults: []int{2, 4},
+		},
+		{
+			input:      FromSlice([]int{}),
+			f:          func(i int) bool { return i%2 == 0 },
+			expResults: nil,
+		},
+		{
+			input:      FromSlice([]int{1, 3, 5, 100}),
+			f:          func(i int) bool { return i > 2 },
+			expResults: []int{3, 5, 100},
+		},
+	} {
+		t.Run(fmt.Sprintf("%v", c.input), func(t *testing.T) {
+			iter := Filter(c.input, c.f)
+			var res []int
+			for iter.Next() {
+				res = append(res, iter.Val())
+			}
+			assert.Equal(t, c.expResults, res)
+		})
+	}
+}

routing/http/types/record_peer.go

@@ -79,3 +79,12 @@ func (pr PeerRecord) MarshalJSON() ([]byte, error) {
 
 	return drjson.MarshalJSONBytes(m)
 }
+
+func FromBitswapRecord(br *BitswapRecord) *PeerRecord {
+	return &PeerRecord{
+		Schema:    SchemaPeer,
+		ID:        br.ID,
+		Addrs:     br.Addrs,
+		Protocols: []string{br.Protocol},
+	}
+}