github.com/zignig/go-ipfs@v0.0.0-20141111235910-c9e5fdf55a52/routing/mock/routing.go (about)

     1  package mock
     2  
     3  import (
     4  	"errors"
     5  	"math/rand"
     6  	"sync"
     7  
     8  	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
     9  	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"
    10  	peer "github.com/jbenet/go-ipfs/peer"
    11  	routing "github.com/jbenet/go-ipfs/routing"
    12  	u "github.com/jbenet/go-ipfs/util"
    13  )
    14  
    15  var _ routing.IpfsRouting = &MockRouter{}
    16  
    17  type MockRouter struct {
    18  	datastore ds.Datastore
    19  	hashTable RoutingServer
    20  	peer      peer.Peer
    21  }
    22  
    23  func NewMockRouter(local peer.Peer, dstore ds.Datastore) routing.IpfsRouting {
    24  	return &MockRouter{
    25  		datastore: dstore,
    26  		peer:      local,
    27  		hashTable: VirtualRoutingServer(),
    28  	}
    29  }
    30  
    31  func (mr *MockRouter) SetRoutingServer(rs RoutingServer) {
    32  	mr.hashTable = rs
    33  }
    34  
    35  func (mr *MockRouter) PutValue(ctx context.Context, key u.Key, val []byte) error {
    36  	return mr.datastore.Put(key.DsKey(), val)
    37  }
    38  
    39  func (mr *MockRouter) GetValue(ctx context.Context, key u.Key) ([]byte, error) {
    40  	v, err := mr.datastore.Get(key.DsKey())
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  
    45  	data, ok := v.([]byte)
    46  	if !ok {
    47  		return nil, errors.New("could not cast value from datastore")
    48  	}
    49  
    50  	return data, nil
    51  }
    52  
    53  func (mr *MockRouter) FindProviders(ctx context.Context, key u.Key) ([]peer.Peer, error) {
    54  	return nil, nil
    55  }
    56  
    57  func (mr *MockRouter) FindPeer(ctx context.Context, pid peer.ID) (peer.Peer, error) {
    58  	return nil, nil
    59  }
    60  
    61  func (mr *MockRouter) FindProvidersAsync(ctx context.Context, k u.Key, max int) <-chan peer.Peer {
    62  	out := make(chan peer.Peer)
    63  	go func() {
    64  		defer close(out)
    65  		for i, p := range mr.hashTable.Providers(k) {
    66  			if max <= i {
    67  				return
    68  			}
    69  			select {
    70  			case out <- p:
    71  			case <-ctx.Done():
    72  				return
    73  			}
    74  		}
    75  	}()
    76  	return out
    77  }
    78  
    79  func (mr *MockRouter) Provide(_ context.Context, key u.Key) error {
    80  	return mr.hashTable.Announce(mr.peer, key)
    81  }
    82  
    83  type RoutingServer interface {
    84  	Announce(peer.Peer, u.Key) error
    85  
    86  	Providers(u.Key) []peer.Peer
    87  
    88  	Client(p peer.Peer) routing.IpfsRouting
    89  }
    90  
    91  func VirtualRoutingServer() RoutingServer {
    92  	return &hashTable{
    93  		providers: make(map[u.Key]peer.Map),
    94  	}
    95  }
    96  
    97  type hashTable struct {
    98  	lock      sync.RWMutex
    99  	providers map[u.Key]peer.Map
   100  }
   101  
   102  func (rs *hashTable) Announce(p peer.Peer, k u.Key) error {
   103  	rs.lock.Lock()
   104  	defer rs.lock.Unlock()
   105  
   106  	_, ok := rs.providers[k]
   107  	if !ok {
   108  		rs.providers[k] = make(peer.Map)
   109  	}
   110  	rs.providers[k][p.Key()] = p
   111  	return nil
   112  }
   113  
   114  func (rs *hashTable) Providers(k u.Key) []peer.Peer {
   115  	rs.lock.RLock()
   116  	defer rs.lock.RUnlock()
   117  	ret := make([]peer.Peer, 0)
   118  	peerset, ok := rs.providers[k]
   119  	if !ok {
   120  		return ret
   121  	}
   122  	for _, peer := range peerset {
   123  		ret = append(ret, peer)
   124  	}
   125  
   126  	for i := range ret {
   127  		j := rand.Intn(i + 1)
   128  		ret[i], ret[j] = ret[j], ret[i]
   129  	}
   130  
   131  	return ret
   132  }
   133  
   134  func (rs *hashTable) Client(p peer.Peer) routing.IpfsRouting {
   135  	return &MockRouter{
   136  		peer:      p,
   137  		hashTable: rs,
   138  	}
   139  }