@@ -47,6 +47,7 @@ func init() {
4747 FindValue : 5 * time .Second ,
4848 BatchFindValues : 60 * time .Second , // responder compresses
4949 BatchGetValues : 75 * time .Second , // large, sometimes cloud fetch then send back
50+ BatchProbeKeys : 10 * time .Second ,
5051 StoreData : 10 * time .Second ,
5152 BatchStoreData : 75 * time .Second , // large uncompressed payloads
5253 Replicate : 90 * time .Second ,
@@ -235,6 +236,37 @@ func (s *Network) handleFindValue(ctx context.Context, message *Message) (res []
235236 return s .encodeMesage (resMsg )
236237}
237238
239+ func (s * Network ) handleBatchProbeKeys (ctx context.Context , message * Message ) (res []byte , err error ) {
240+ defer func () {
241+ if response , err := s .handlePanic (ctx , message .Sender , BatchProbeKeys ); response != nil || err != nil {
242+ res = response
243+ }
244+ }()
245+
246+ request , ok := message .Data .(* BatchProbeKeysRequest )
247+ if ! ok {
248+ err := errors .New ("invalid BatchProbeKeysRequest" )
249+ return s .generateResponseMessage (ctx , BatchProbeKeys , message .Sender , ResultFailed , err .Error ())
250+ }
251+
252+ // Keep routing table fresh while handling probe traffic.
253+ s .dht .addNode (ctx , message .Sender )
254+
255+ statuses , err := s .dht .store .RetrieveBatchLocalStatus (ctx , request .Keys )
256+ if err != nil {
257+ err = errors .Errorf ("batch probe keys: %w" , err )
258+ return s .generateResponseMessage (ctx , BatchProbeKeys , message .Sender , ResultFailed , err .Error ())
259+ }
260+
261+ response := & BatchProbeKeysResponse {
262+ Status : ResponseStatus {Result : ResultOk },
263+ Data : statuses ,
264+ }
265+
266+ resMsg := s .dht .newMessage (BatchProbeKeys , message .Sender , response )
267+ return s .encodeMesage (resMsg )
268+ }
269+
238270func (s * Network ) handleStoreData (ctx context.Context , message * Message ) (res []byte , err error ) {
239271 defer func () {
240272 if response , err := s .handlePanic (ctx , message .Sender , StoreData ); response != nil || err != nil {
@@ -480,6 +512,10 @@ func (s *Network) handleConn(ctx context.Context, rawConn net.Conn) {
480512 response , hErr = s .withMetrics (BatchGetValues , func () ([]byte , error ) {
481513 return s .handleGetValuesRequest (ctx , request , reqID )
482514 })
515+ case BatchProbeKeys :
516+ response , hErr = s .withMetrics (BatchProbeKeys , func () ([]byte , error ) {
517+ return s .handleBatchProbeKeys (ctx , request )
518+ })
483519 default :
484520 // count unknown types as failure and return
485521 m := s .metricsFor (mt )
@@ -1321,6 +1357,8 @@ func (s *Network) generateResponseMessage(ctx context.Context, messageType int,
13211357 response = & ReplicateDataResponse {Status : responseStatus }
13221358 case BatchGetValues :
13231359 response = & BatchGetValuesResponse {Status : responseStatus }
1360+ case BatchProbeKeys :
1361+ response = & BatchProbeKeysResponse {Status : responseStatus }
13241362 default :
13251363 return nil , fmt .Errorf ("unsupported message type %d" , messageType )
13261364 }
@@ -1461,6 +1499,8 @@ func msgName(t int) string {
14611499 return "BatchStoreData"
14621500 case Replicate :
14631501 return "Replicate"
1502+ case BatchProbeKeys :
1503+ return "BatchProbeKeys"
14641504 default :
14651505 return fmt .Sprintf ("Type_%d" , t )
14661506 }
@@ -1501,7 +1541,7 @@ func (s *Network) HandleMetricsSnapshot() map[string]HandleCounters {
15011541func readDeadlineFor (msgType int , overall time.Duration ) time.Duration {
15021542 small := 10 * time .Second
15031543 switch msgType {
1504- case Ping , FindNode , BatchFindNode , FindValue , StoreData , BatchStoreData :
1544+ case Ping , FindNode , BatchFindNode , FindValue , StoreData , BatchStoreData , BatchProbeKeys :
15051545 if overall > small + 1 * time .Second {
15061546 return small
15071547 }
0 commit comments