more stuff

rpc-pool.go

@@ -1,13 +1,14 @@
 package main
 
 import (
+	"fmt"
 	"io"
 	"net"
 	"strings"
 )
 
 type RemoteGrainPool struct {
-	Hosts  []string
+	Host   string
 	grains map[CartId]RemoteGrain
 }
 
@@ -46,7 +47,6 @@ func (g *RemoteGrain) Connect() error {
 }
 
 func (g *RemoteGrain) HandleMessage(message *Message, isReplay bool) ([]byte, error) {
-
 	err := SendCartPacket(g.client, g.Id, RemoteHandleMessage, message.Write)
 	if err != nil {
 		return nil, err
@@ -60,7 +60,6 @@ func (g *RemoteGrain) GetId() CartId {
 }
 
 func (g *RemoteGrain) GetCurrentState() ([]byte, error) {
-
 	err := SendCartPacket(g.client, g.Id, RemoteGetState, func(w io.Writer) error {
 		return nil
 	})
@@ -71,9 +70,9 @@ func (g *RemoteGrain) GetCurrentState() ([]byte, error) {
 	return data, err
 }
 
-func NewRemoteGrainPool(addr ...string) *RemoteGrainPool {
+func NewRemoteGrainPool(addr string) *RemoteGrainPool {
 	return &RemoteGrainPool{
-		Hosts:  addr,
+		Host:   addr,
 		grains: make(map[CartId]RemoteGrain),
 	}
 }
@@ -83,17 +82,26 @@ func (p *RemoteGrainPool) findRemoteGrain(id CartId) *RemoteGrain {
 	if !ok {
 		return nil
 	}
+	grain.Connect()
 	return &grain
 }
 
+func (p *RemoteGrainPool) findOrCreateGrain(id CartId) *RemoteGrain {
+	grain := p.findRemoteGrain(id)
+	if grain == nil {
+		grain = NewRemoteGrain(id, p.Host)
+		p.grains[id] = *grain
+		grain.Connect()
+	}
+	return grain
+}
+
 func (p *RemoteGrainPool) Process(id CartId, messages ...Message) ([]byte, error) {
 	var result []byte
 	var err error
-	grain := p.findRemoteGrain(id)
+	grain := p.findOrCreateGrain(id)
 	if grain == nil {
-		grain = NewRemoteGrain(id, p.Hosts[0])
-		grain.Connect()
-		p.grains[id] = *grain
+		return nil, fmt.Errorf("grain not found")
 	}
 	for _, message := range messages {
 		result, err = grain.HandleMessage(&message, false)
@@ -102,9 +110,9 @@ func (p *RemoteGrainPool) Process(id CartId, messages ...Message) ([]byte, error
 }
 
 func (p *RemoteGrainPool) Get(id CartId) ([]byte, error) {
-	grain := p.findRemoteGrain(id)
+	grain := p.findOrCreateGrain(id)
 	if grain == nil {
-		return nil, nil
+		return nil, fmt.Errorf("grain not found")
 	}
 	return grain.GetCurrentState()
 }