LBRYCommunity/lbcd
1
0
Fork 0
Code Issues 17 Pull requests 4 Projects Releases 28 Packages Wiki Activity

Stop rescans when a websocket client disconnects.

Browse source 
Fixes #66.
This commit is contained in:
Josh Rickmar 2014-01-14 15:59:31 -05:00
parent 6abad1d8ac
commit 051a9013ce
1 changed files with 136 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

152
rpcwebsocket.go
View file

@ -22,7 +22,12 @@ import (
type ntfnChan chan btcjson.Cmd
type wsCommandHandler func(*rpcServer, btcjson.Cmd, ntfnChan) (interface{}, *btcjson.Error)
type handlerChans struct {
n ntfnChan // channel to send notifications
disconnected <-chan struct{} // closed when a client has disconnected.
}
type wsCommandHandler func(*rpcServer, btcjson.Cmd, handlerChans) (interface{}, *btcjson.Error)
// wsHandlers maps RPC command strings to appropriate websocket handler
// functions.
@ -226,7 +231,7 @@ type requestContexts struct {
// extension JSON-RPC command and runs the proper handler to reply to
// the command. Any and all responses are sent to the wallet from
// this function.
func respondToAnyCmd(cmd btcjson.Cmd, s *rpcServer, n ntfnChan) *btcjson.Reply {
func respondToAnyCmd(cmd btcjson.Cmd, s *rpcServer, c handlerChans) *btcjson.Reply {
// Lookup the websocket extension for the command and if it doesn't
// exist fallback to handling the command as a standard command.
wsHandler, ok := wsHandlers[cmd.Method()]
@ -236,7 +241,7 @@ func respondToAnyCmd(cmd btcjson.Cmd, s *rpcServer, n ntfnChan) *btcjson.Reply {
response := standardCmdReply(cmd, s)
return &response
}
result, jsonErr := wsHandler(s, cmd, n)
result, jsonErr := wsHandler(s, cmd, c)
id := cmd.Id()
response := btcjson.Reply{
Id: &id,
@ -248,7 +253,7 @@ func respondToAnyCmd(cmd btcjson.Cmd, s *rpcServer, n ntfnChan) *btcjson.Reply {
// handleGetCurrentNet implements the getcurrentnet command extension
// for websocket connections.
func handleGetCurrentNet(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{}, *btcjson.Error) {
func handleGetCurrentNet(s *rpcServer, icmd btcjson.Cmd, c handlerChans) (interface{}, *btcjson.Error) {
if cfg.TestNet3 {
return btcwire.TestNet3, nil
}
@ -257,7 +262,7 @@ func handleGetCurrentNet(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interf
// handleGetBestBlock implements the getbestblock command extension
// for websocket connections.
func handleGetBestBlock(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{}, *btcjson.Error) {
func handleGetBestBlock(s *rpcServer, icmd btcjson.Cmd, c handlerChans) (interface{}, *btcjson.Error) {
// All other "get block" commands give either the height, the
// hash, or both but require the block SHA. This gets both for
// the best block.
@ -276,7 +281,7 @@ func handleGetBestBlock(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interfa
// handleNotifyNewTXs implements the notifynewtxs command extension for
// websocket connections.
func handleNotifyNewTXs(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{}, *btcjson.Error) {
func handleNotifyNewTXs(s *rpcServer, icmd btcjson.Cmd, c handlerChans) (interface{}, *btcjson.Error) {
cmd, ok := icmd.(*btcws.NotifyNewTXsCmd)
if !ok {
return nil, &btcjson.ErrInternal
@ -302,7 +307,7 @@ func handleNotifyNewTXs(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interfa
return nil, &e
}
s.ws.AddTxRequest(ntfns, addr.EncodeAddress())
s.ws.AddTxRequest(c.n, addr.EncodeAddress())
}
return nil, nil
@ -310,20 +315,20 @@ func handleNotifyNewTXs(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interfa
// handleNotifySpent implements the notifyspent command extension for
// websocket connections.
func handleNotifySpent(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{}, *btcjson.Error) {
func handleNotifySpent(s *rpcServer, icmd btcjson.Cmd, c handlerChans) (interface{}, *btcjson.Error) {
cmd, ok := icmd.(*btcws.NotifySpentCmd)
if !ok {
return nil, &btcjson.ErrInternal
}
s.ws.AddSpentRequest(ntfns, cmd.OutPoint)
s.ws.AddSpentRequest(c.n, cmd.OutPoint)
return nil, nil
}
// handleRescan implements the rescan command extension for websocket
// connections.
func handleRescan(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{}, *btcjson.Error) {
func handleRescan(s *rpcServer, icmd btcjson.Cmd, c handlerChans) (interface{}, *btcjson.Error) {
cmd, ok := icmd.(*btcws.RescanCmd)
if !ok {
return nil, &btcjson.ErrInternal
@ -357,6 +362,34 @@ func handleRescan(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{},
rpcsLog.Errorf("Error looking up block sha: %v", err)
return nil, &btcjson.ErrDatabase
}
// A select statement is used to stop rescans if the
// client requesting the rescan has disconnected.
select {
case <-c.disconnected:
rpcsLog.Infof("Stopping rescan at height %v for disconnected client",
blk.Height())
return nil, nil
default:
rescanBlock(s, cmd, c, blk)
}
}
if maxblock-minblock > int64(len(blkshalist)) {
minblock += int64(len(blkshalist))
} else {
break
}
}
rpcsLog.Info("Finished rescan")
return nil, nil
}
// rescanBlock rescans all transactions in a single block. This is a
// helper function for handleRescan.
func rescanBlock(s *rpcServer, cmd *btcws.RescanCmd, c handlerChans, blk *btcutil.Block) {
for _, tx := range blk.Transactions() {
var txReply *btcdb.TxListReply
txouts:
@ -367,9 +400,11 @@ func handleRescan(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{},
continue txouts
}
for i, addr := range addrs {
for _, addr := range addrs {
encodedAddr := addr.EncodeAddress()
if _, ok := cmd.Addresses[encodedAddr]; ok {
if _, ok := cmd.Addresses[encodedAddr]; !ok {
continue
}
// TODO(jrick): This lookup is expensive and can be avoided
// if the wallet is sent the previous outpoints for all inputs
// of the tx, so any can removed from the utxo set (since
@ -389,39 +424,37 @@ func handleRescan(s *rpcServer, icmd btcjson.Cmd, ntfns ntfnChan) (interface{},
}
n := &btcws.ProcessedTxNtfn{
// Sha never errors.
blksha, _ := blk.Sha()
ntfn := &btcws.ProcessedTxNtfn{
Receiver: encodedAddr,
Amount: txout.Value,
TxID: tx.Sha().String(),
TxOutIndex: uint32(txOutIdx),
PkScript: hex.EncodeToString(txout.PkScript),
BlockHash: blkshalist[i].String(),
BlockHash: blksha.String(),
BlockHeight: int32(blk.Height()),
BlockIndex: tx.Index(),
BlockTime: blk.MsgBlock().Header.Timestamp.Unix(),
Spent: txReply.TxSpent[txOutIdx],
}
ntfns <- n
}
}
}
}
}
if maxblock-minblock > int64(len(blkshalist)) {
minblock += int64(len(blkshalist))
} else {
break
}
}
select {
case <-c.disconnected:
return
rpcsLog.Info("Finished rescan")
return nil, nil
default:
c.n <- ntfn
}
}
}
}
}
// handleWalletSendRawTransaction implements the websocket extended version of
// the sendrawtransaction command.
func handleWalletSendRawTransaction(s *rpcServer, icmd btcjson.Cmd, n ntfnChan) (interface{}, *btcjson.Error) {
func handleWalletSendRawTransaction(s *rpcServer, icmd btcjson.Cmd, c handlerChans) (interface{}, *btcjson.Error) {
result, err := handleSendRawTransaction(s, icmd)
// TODO: the standard handlers really should be changed to
// return btcjson.Errors which get used directly in the
@ -443,7 +476,7 @@ func handleWalletSendRawTransaction(s *rpcServer, icmd btcjson.Cmd, n ntfnChan)
txSha, _ := btcwire.NewShaHashFromStr(result.(string))
// Request to be notified when the transaction is mined.
s.ws.AddMinedTxRequest(n, txSha)
s.ws.AddMinedTxRequest(c.n, txSha)
return result, nil
}
@ -513,6 +546,13 @@ func (s *rpcServer) walletReqsNotifications(ws *websocket.Conn) {
// Channel for responses.
r := make(chan *btcjson.Reply)
// Channels for websocket handlers.
disconnected := make(chan struct{})
hc := handlerChans{
n: n,
disconnected: disconnected,
}
// msgs is a channel for all messages received over the websocket.
msgs := make(chan []byte)
@ -522,13 +562,12 @@ func (s *rpcServer) walletReqsNotifications(ws *websocket.Conn) {
for {
select {
case <-s.quit:
close(msgs)
return
default:
var m []byte
if err := websocket.Message.Receive(ws, &m); err != nil {
close(msgs)
close(disconnected)
return
}
msgs <- m
@ -538,13 +577,27 @@ func (s *rpcServer) walletReqsNotifications(ws *websocket.Conn) {
for {
select {
case m, ok := <-msgs:
if !ok {
// Wallet disconnected.
case <-s.quit:
// Server closed. Closing disconnected signals handlers to stop
// and flushes all channels handlers may write to.
close(disconnected)
case <-disconnected:
for {
select {
case <-msgs:
case <-r:
case <-n:
default:
return
}
// Handle request here.
go s.websocketJSONHandler(r, n, m)
}
case m := <-msgs:
// Spawn new goroutine to handle request. Responses and
// notifications are read by channels in this for-select
// loop.
go s.websocketJSONHandler(r, hc, m)
case response := <-r:
// Marshal and send response.
@ -555,6 +608,7 @@ func (s *rpcServer) walletReqsNotifications(ws *websocket.Conn) {
}
if err := websocket.Message.Send(ws, mresp); err != nil {
// Wallet disconnected.
close(disconnected)
return
}
@ -567,25 +621,24 @@ func (s *rpcServer) walletReqsNotifications(ws *websocket.Conn) {
}
if err := websocket.Message.Send(ws, mntfn); err != nil {
// Wallet disconnected.
close(disconnected)
return
}
case <-s.quit:
// Server closed.
return
}
}
}
// websocketJSONHandler parses and handles a marshalled json message,
// sending the marshalled reply to a wallet notification channel.
func (s *rpcServer) websocketJSONHandler(r chan *btcjson.Reply, n ntfnChan, msg []byte) {
func (s *rpcServer) websocketJSONHandler(r chan *btcjson.Reply, c handlerChans, msg []byte) {
s.wg.Add(1)
defer s.wg.Done()
var resp *btcjson.Reply
cmd, jsonErr := parseCmd(msg)
if jsonErr != nil {
var resp btcjson.Reply
resp = &btcjson.Reply{}
if cmd != nil {
// Unmarshaling at least a valid JSON-RPC message succeeded.
// Use the provided id for errors. Requests with no IDs
@ -597,13 +650,20 @@ func (s *rpcServer) websocketJSONHandler(r chan *btcjson.Reply, n ntfnChan, msg
resp.Id = &id
}
resp.Error = jsonErr
r <- &resp
return
} else {
resp = respondToAnyCmd(cmd, s, c)
}
resp := respondToAnyCmd(cmd, s, n)
// Once response has been processed, only send if the client
// is still connected.
select {
case <-c.disconnected:
return
default:
r <- resp
}
}
// NotifyBlockConnected creates and marshalls a JSON message to notify
// of a new block connected to the main chain. The notification is sent