diff --git a/bchain/mempool.go b/bchain/mempool.go index 558a514c..9aa4cd79 100644 --- a/bchain/mempool.go +++ b/bchain/mempool.go @@ -8,12 +8,28 @@ import ( "github.com/golang/glog" ) +type scriptIndex struct { + script string + n uint32 +} + +type outpoint struct { + txid string + vout uint32 +} + +type inputOutput struct { + outputScripts []scriptIndex + inputs []outpoint +} + // Mempool is mempool handle. type Mempool struct { - chain *BitcoinRPC - mux sync.Mutex - scriptToTx map[string][]string - txToScript map[string][]string + chain *BitcoinRPC + mux sync.Mutex + txToInputOutput map[string]inputOutput + scriptToTx map[string][]outpoint + inputs map[outpoint]string } // NewMempool creates new mempool handler. @@ -26,14 +42,30 @@ func (m *Mempool) GetTransactions(outputScript []byte) ([]string, error) { m.mux.Lock() defer m.mux.Unlock() scriptHex := hex.EncodeToString(outputScript) - return m.scriptToTx[scriptHex], nil + outpoints := m.scriptToTx[scriptHex] + txs := make([]string, 0, len(outpoints)+len(outpoints)/2) + for _, o := range outpoints { + txs = append(txs, o.txid) + i := m.inputs[o] + if i != "" { + txs = append(txs, i) + } + } + return txs, nil } -func (m *Mempool) updateMaps(newScriptToTx map[string][]string, newTxToScript map[string][]string) { +// GetInput returns transaction which spends given outpoint +func (m *Mempool) GetInput(outputTxid string, vout uint32) string { + o := outpoint{txid: outputTxid, vout: vout} + return m.inputs[o] +} + +func (m *Mempool) updateMappings(newTxToInputOutput map[string]inputOutput, newScriptToTx map[string][]outpoint, newInputs map[outpoint]string) { m.mux.Lock() defer m.mux.Unlock() + m.txToInputOutput = newTxToInputOutput m.scriptToTx = newScriptToTx - m.txToScript = newTxToScript + m.inputs = newInputs } // Resync gets mempool transactions and maps output scripts to transactions. @@ -46,30 +78,41 @@ func (m *Mempool) Resync() error { if err != nil { return err } - newScriptToTx := make(map[string][]string) - newTxToScript := make(map[string][]string) + newTxToInputOutput := make(map[string]inputOutput, len(m.txToInputOutput)+1) + newScriptToTx := make(map[string][]outpoint, len(m.scriptToTx)+1) + newInputs := make(map[outpoint]string, len(m.inputs)+1) for _, txid := range txs { - scripts := m.txToScript[txid] - if scripts == nil { + io, exists := m.txToInputOutput[txid] + if !exists { tx, err := m.chain.GetTransaction(txid) if err != nil { glog.Error("cannot get transaction ", txid, ": ", err) continue } - scripts = make([]string, 0, len(tx.Vout)) + io.outputScripts = make([]scriptIndex, 0, len(tx.Vout)) for _, output := range tx.Vout { outputScript := output.ScriptPubKey.Hex if outputScript != "" { - scripts = append(scripts, outputScript) + io.outputScripts = append(io.outputScripts, scriptIndex{outputScript, output.N}) } } + io.inputs = make([]outpoint, 0, len(tx.Vin)) + for _, input := range tx.Vin { + if input.Coinbase != "" { + continue + } + io.inputs = append(io.inputs, outpoint{input.Txid, input.Vout}) + } } - newTxToScript[txid] = scripts - for _, script := range scripts { - newScriptToTx[script] = append(newScriptToTx[script], txid) + newTxToInputOutput[txid] = io + for _, si := range io.outputScripts { + newScriptToTx[si.script] = append(newScriptToTx[si.script], outpoint{txid, si.n}) + } + for _, i := range io.inputs { + newInputs[i] = txid } } - m.updateMaps(newScriptToTx, newTxToScript) - glog.Info("Mempool: resync finished in ", time.Since(start), ", ", len(m.txToScript), " transactions in mempool") + m.updateMappings(newTxToInputOutput, newScriptToTx, newInputs) + glog.Info("Mempool: resync finished in ", time.Since(start), ", ", len(m.txToInputOutput), " transactions in mempool") return nil }