commit f80086aad0652c4386d0facc6432d8065991e482
parent 8a5d996c1294545c35a9d2fe6c313184ba415d31
Author: MTRNord <mtrnord1@gmail.com>
Date: Sat, 16 Mar 2024 23:11:38 +0100
Improve performance and memory requirements of the map wrapper by not doing as many copies
Diffstat:
6 files changed, 99 insertions(+), 67 deletions(-)
diff --git a/.gitignore b/.gitignore
@@ -2,4 +2,6 @@ descriptor.pb
docs
matrix_protobuf_fed
bin/
-.mypy_cache
-\ No newline at end of file
+.mypy_cache
+*.prof
+*.mprof
+\ No newline at end of file
diff --git a/cmds/client/main.go b/cmds/client/main.go
@@ -37,18 +37,25 @@ func main() {
log.Println("Created client...")
timeoutCtx, cancel := context.WithTimeout(context.TODO(), time.Second*10)
+ defer cancel()
err = printServerVersion(timeoutCtx, client)
if err != nil {
log.Println("Failed to print server version:", err)
}
- cancel()
timeoutCtx, cancel = context.WithTimeout(context.TODO(), time.Second*10)
+ defer cancel()
err = printServerKeys(timeoutCtx, client)
if err != nil {
log.Println("Failed to print server version:", err)
}
- cancel()
+
+ timeoutCtx, cancel = context.WithTimeout(context.TODO(), time.Second*10)
+ defer cancel()
+ err = printServerVersion(timeoutCtx, client)
+ if err != nil {
+ log.Println("Failed to print server version:", err)
+ }
}
func printServerKeys(ctx context.Context, client protocol.MatrixFederation) error {
diff --git a/cmds/server/main.go b/cmds/server/main.go
@@ -3,8 +3,15 @@ package main
import (
"context"
"errors"
+ "flag"
+ "fmt"
"log"
+ "log/slog"
"net"
+ "os"
+ "os/signal"
+ "runtime/pprof"
+ "syscall"
"capnproto.org/go/capnp/v3"
"capnproto.org/go/capnp/v3/flowcontrol"
@@ -32,17 +39,22 @@ func Serve(lis net.Listener, boot capnp.Client) error {
// Accept incoming connections
conn, err := lis.Accept()
if err != nil {
+ conn.Close()
return err
}
+ defer conn.Close()
// the RPC connection takes ownership of the bootstrap interface and will release it when the connection
// exits, so use AddRef to avoid releasing the provided bootstrap client capability.
opts := rpc.Options{
BootstrapClient: boot.AddRef(),
+ Logger: slog.Default(),
}
// For each new incoming connection, create a new RPC transport connection that will serve incoming RPC requests
transport := rpc.NewStreamTransport(conn)
- _ = rpc.NewConn(transport, &opts)
+ defer transport.Close()
+ rpc_conn := rpc.NewConn(transport, &opts)
+ defer rpc_conn.Close()
}
}
@@ -66,7 +78,44 @@ func ListenAndServe(ctx context.Context, network, addr string, bootstrapClient c
return err
}
+var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
+var memprofile = flag.String("memprofile", "", "write memory profile to this file")
+var f *os.File
+
func main() {
+ flag.Parse()
+ if *cpuprofile != "" {
+ f, err := os.Create(*cpuprofile)
+ if err != nil {
+ log.Fatal(err)
+ }
+ pprof.StartCPUProfile(f)
+ defer pprof.StopCPUProfile()
+ }
+
+ c := make(chan os.Signal, 2)
+ signal.Notify(c, os.Interrupt, syscall.SIGTERM) // subscribe to system signals
+ onKill := func(c chan os.Signal) {
+ select {
+ case <-c:
+ fmt.Println("Got killed")
+ pprof.StopCPUProfile()
+ f.Close()
+ if *memprofile != "" {
+ f, err := os.Create(*memprofile)
+ if err != nil {
+ log.Fatal(err)
+ }
+ pprof.WriteHeapProfile(f)
+ f.Close()
+ }
+ os.Exit(0)
+ }
+ }
+
+ // try to handle os interrupt(signal terminated)
+ go onKill(c)
+
log.Println("Starting server on port localhost:8449")
server := rpcserver.NewServer()
diff --git a/rpcserver/map_wrapper.go b/rpcserver/map_wrapper.go
@@ -22,47 +22,50 @@ import (
* Contrary to normal go maps this map has a fixed size.
*/
type Map[Key capnp.Ptr, Value capnp.Ptr] struct {
- internal_map *protocol_types.Map
+ internalMap *protocol_types.Map
+ maxSize int32
}
// NewMap creates a new Map Wrapper
-func NewMap[Key capnp.Ptr, Value capnp.Ptr](s *capnp.Segment) (*Map[Key, Value], error) {
- internal_map, err := protocol_types.NewMap(s)
+func NewMap[Key capnp.Ptr, Value capnp.Ptr](s *capnp.Segment, maxSize int32) (*Map[Key, Value], error) {
+ internalMap, err := protocol_types.NewMap(s)
if err != nil {
return nil, err
}
return &Map[Key, Value]{
- internal_map: &internal_map,
+ internalMap: &internalMap,
+ maxSize: maxSize,
}, nil
}
// FromMap converts a capnp map to a wrapper
-func FromMap[Key capnp.Ptr, Value capnp.Ptr](m *protocol_types.Map) *Map[Key, Value] {
+func FromMap[Key capnp.Ptr, Value capnp.Ptr](m *protocol_types.Map, maxSize int32) *Map[Key, Value] {
return &Map[Key, Value]{
- internal_map: m,
+ internalMap: m,
+ maxSize: maxSize,
}
}
// HasEntries returns true if the map has entries
func (m *Map[Key, Value]) HasEntries() bool {
- return m.internal_map.HasEntries()
+ return m.internalMap.HasEntries()
}
// Get a Segment of the internal map
func (m *Map[Key, Value]) Segment() *capnp.Segment {
- return m.internal_map.Segment()
+ return m.internalMap.Segment()
}
// Entries returns the entries of the map as a go map
-func (m *Map[Key, Value]) Entries() (map[Key]Value, error) {
+func (m *Map[Key, Value]) Entries() (map[Key]*Value, error) {
// Check if we have entries. If not we return an empty map
- result := make(map[Key]Value)
- if !m.internal_map.HasEntries() {
+ result := make(map[Key]*Value)
+ if !m.HasEntries() {
return result, nil
}
- entries, err := m.internal_map.Entries()
+ entries, err := m.internalMap.Entries()
if err != nil {
return nil, err
}
@@ -74,51 +77,36 @@ func (m *Map[Key, Value]) Entries() (map[Key]Value, error) {
return nil, err
}
- value, err := entry.Value()
+ value_raw, err := entry.Value()
if err != nil {
return nil, err
}
- result[Key(key)] = Value(value)
+ value := Value(value_raw)
+ result[Key(key)] = &value
}
return result, nil
}
-type ErrMapTooLarge struct{}
-
-func (e ErrMapTooLarge) Error() string {
- return "Map supplied is larger than the internal map."
-}
-
-// SetEntries sets the entries of the map
-func (m *Map[Key, Value]) SetEntries(entries map[Key]Value) error {
+func (m *Map[Key, Value]) AddEntry(key Key, value Value) error {
// Check if we have any entries
- if !m.internal_map.HasEntries() {
+ if !m.internalMap.HasEntries() {
// Allocate enough entries
- _, err := m.internal_map.NewEntries(int32(len(entries)))
+ _, err := m.internalMap.NewEntries(m.maxSize)
if err != nil {
return err
}
}
-
- // Ensure the map is not larger than the internal map
- internal_entries, err := m.internal_map.Entries()
+ internalEntries, err := m.internalMap.Entries()
if err != nil {
return err
}
- if len(entries) > internal_entries.Len() {
- return ErrMapTooLarge{}
- }
- // Set the entries. Important: We cant use the Entries() method we defined earlier in this struct as that one is a copy.
- idx := 0
- for key, value := range entries {
- entry := internal_entries.At(idx)
- entry.SetKey(capnp.Ptr(key))
- entry.SetValue(capnp.Ptr(value))
- idx++
+ entry := internalEntries.At(internalEntries.Len() - 1)
+ err = entry.SetKey(capnp.Ptr(key))
+ if err != nil {
+ return err
}
-
- return nil
+ return entry.SetValue(capnp.Ptr(value))
}
diff --git a/rpcserver/rpc_server.go b/rpcserver/rpc_server.go
@@ -50,6 +50,7 @@ func NewServer() RPCMatrixServer {
}
func (s RPCMatrixServer) GetVersion(ctx context.Context, call protocol.MatrixFederation_getVersion) error {
+ call.Go()
res, err := call.AllocResults() // Allocate the results struct
if err != nil {
return err
@@ -66,14 +67,12 @@ func (s RPCMatrixServer) GetVersion(ctx context.Context, call protocol.MatrixFed
}
func (s RPCMatrixServer) GetKeys(ctx context.Context, call protocol.MatrixFederation_getKeys) error {
- _, err := call.AllocResults() // Allocate the results struct
- if err != nil {
- return err
- }
+ call.Go()
client := call.Args().Callback()
+ defer client.Release()
- err = client.Write(ctx, func(p protocol.StreamCallback_write_Params) error {
+ err := client.Write(ctx, func(p protocol.StreamCallback_write_Params) error {
log.Println("Sending server keys metadata response...")
response, err := types.NewServerKeysResponse(p.Segment())
@@ -120,11 +119,7 @@ func (s RPCMatrixServer) GetKeys(ctx context.Context, call protocol.MatrixFedera
if err != nil {
return err
}
- verify_keys := FromMap(&verify_keys_raw)
- verify_keys_entries, err := verify_keys.Entries()
- if err != nil {
- return err
- }
+ verify_keys := FromMap(&verify_keys_raw, 1)
key, err := capnp.NewText(verify_keys.Segment(), "placeholder")
if err != nil {
return err
@@ -135,9 +130,7 @@ func (s RPCMatrixServer) GetKeys(ctx context.Context, call protocol.MatrixFedera
if err != nil {
return err
}
-
- verify_keys_entries[key.ToPtr()] = data.ToPtr()
- verify_keys.SetEntries(verify_keys_entries)
+ verify_keys.AddEntry(key.ToPtr(), data.ToPtr())
verify_keys_bytes, err := capnp.Canonicalize(capnp.Struct(verify_keys_raw))
if err != nil {
@@ -157,9 +150,8 @@ func (s RPCMatrixServer) GetKeys(ctx context.Context, call protocol.MatrixFedera
return err
}
- future, release := client.Done(ctx, nil)
+ _, release := client.Done(ctx, nil)
defer release()
- _, err = future.Struct()
if err := client.WaitStreaming(); err != nil {
return err
diff --git a/rpcserver/utils.go b/rpcserver/utils.go
@@ -35,12 +35,7 @@ func SignCapnproto(signingName string, keyID KeyID, privateKey ed25519.PrivateKe
if err != nil {
return err
}
- signatures_map_wrapper := FromMap(&signatures_map)
- signatures_map_go, err := signatures_map_wrapper.Entries()
- if err != nil {
- return err
- }
-
+ signatures_map_wrapper := FromMap(&signatures_map, 1)
key, err := capnp.NewText(signatures_map_wrapper.Segment(), string(keyID))
if err != nil {
return err
@@ -50,8 +45,7 @@ func SignCapnproto(signingName string, keyID KeyID, privateKey ed25519.PrivateKe
if err != nil {
return err
}
- signatures_map_go[key.ToPtr()] = data.ToPtr()
- signatures_map_wrapper.SetEntries(signatures_map_go)
+ signatures_map_wrapper.AddEntry(key.ToPtr(), data.ToPtr())
return nil
}
