summaryrefslogtreecommitdiffstats
path: root/vespaclient-core/src/main/java/com/yahoo/feedapi/MessageBusSessionFactory.java
blob: 54e638717e0b701342d7c29b5a4a3552160bb02d (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.feedapi;

import com.yahoo.cloud.config.SlobroksConfig;
import com.yahoo.document.config.DocumentmanagerConfig;
import com.yahoo.documentapi.VisitorParameters;
import com.yahoo.documentapi.VisitorSession;
import com.yahoo.documentapi.messagebus.MessageBusDocumentAccess;
import com.yahoo.documentapi.messagebus.MessageBusParams;
import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
import com.yahoo.documentapi.messagebus.protocol.RemoveDocumentMessage;
import com.yahoo.documentapi.messagebus.protocol.UpdateDocumentMessage;
import com.yahoo.jdisc.Metric;
import com.yahoo.messagebus.Message;
import com.yahoo.messagebus.ReplyHandler;
import com.yahoo.messagebus.SourceSession;
import com.yahoo.messagebus.network.rpc.RPCNetworkParams;

import java.util.Collections;

public class MessageBusSessionFactory implements SessionFactory {

    private final MessageBusDocumentAccess access;
    private final MessagePropertyProcessor processor;

    private interface Metrics {
        String NUM_OPERATIONS = "num_operations";
        String NUM_PUTS = "num_puts";
        String NUM_REMOVES = "num_removes";
        String NUM_UPDATES = "num_updates";
    }

    @SuppressWarnings("unused") // used from extensions
    public MessageBusSessionFactory(MessagePropertyProcessor processor) {
        this(processor, null, null);
    }
    
    public MessageBusSessionFactory(MessagePropertyProcessor processor, 
                                    DocumentmanagerConfig documentmanagerConfig,
                                    SlobroksConfig slobroksConfig) {
        this.processor = processor;
        MessageBusParams params = new MessageBusParams(processor.getLoadTypes());
        params.setTraceLevel(processor.getFeederOptions().getTraceLevel());
        RPCNetworkParams rpcNetworkParams = processor.getFeederOptions().getNetworkParams();
        if (slobroksConfig != null) // not set: will subscribe
            rpcNetworkParams.setSlobroksConfig(slobroksConfig);
        params.setRPCNetworkParams(rpcNetworkParams);
        params.setDocumentManagerConfigId("client");
        if (documentmanagerConfig != null) // not set: will subscribe
            params.setDocumentmanagerConfig(documentmanagerConfig);
        access = new MessageBusDocumentAccess(params);
    }

    public MessageBusDocumentAccess getAccess() {
        return access;
    }

    @Override
    public synchronized SendSession createSendSession(ReplyHandler handler, Metric metric) {
        return new SourceSessionWrapper(
                access.getMessageBus().createSourceSession(handler, processor.getFeederOptions().toSourceSessionParams()),
                metric);
    }

    public void shutDown() {
        access.shutdown();
    }

    @Override
    public synchronized VisitorSession createVisitorSession(VisitorParameters params) {
        try {
            return access.createVisitorSession(params);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private class SourceSessionWrapper extends SendSession {

        private final SourceSession session;
        private final Metric metric;
        private final Metric.Context context;

        private SourceSessionWrapper(SourceSession session, Metric metric) {
            this.session = session;
            this.metric = metric;
            this.context = metric.createContext(Collections.<String, String>emptyMap());
        }

        @Override
        protected com.yahoo.messagebus.Result onSend(Message m, boolean blockIfQueueFull) throws InterruptedException {
            updateCounters(m);
            if (blockIfQueueFull) {
                return session.sendBlocking(m);
            } else {
                return session.send(m);
            }
        }

        private void updateCounters(Message m) {
            metric.add(Metrics.NUM_OPERATIONS, 1, context);

            if (m instanceof PutDocumentMessage) {
                metric.add(Metrics.NUM_PUTS, 1, context);
            } else if (m instanceof RemoveDocumentMessage) {
                metric.add(Metrics.NUM_REMOVES, 1, context);
            } else if (m instanceof UpdateDocumentMessage) {
                metric.add(Metrics.NUM_UPDATES, 1, context);
            }
        }

        @Override
        public void close() {
            session.close();
        }
    }

}